From 31b4a9df6cf7e4a30a76db179b9fceccb1226cd3 Mon Sep 17 00:00:00 2001 From: nihalp Date: Mon, 17 Aug 2026 23:07:22 +0530 Subject: [PATCH 01/99] feat(cuteir): add shared semantic frontend --- Cargo.lock | 30 + Cargo.toml | 8 +- README.md | 3 + crates/cargo-oxide/src/commands/context.rs | 61 +- .../cargo-oxide/src/commands/examples_list.rs | 26 +- crates/cargo-oxide/src/commands/tests.rs | 43 + crates/cargo-oxide/src/main.rs | 4 +- crates/cuda-oxide-codegen/Cargo.toml | 1 + crates/cuda-oxide-codegen/src/pipeline.rs | 9 + crates/cuda-oxide-codegen/src/prep.rs | 572 +++ crates/mir-importer/Cargo.toml | 2 + .../src/translator/cute/block_scaled.rs | 254 + .../src/translator/cute/block_scaled_emit.rs | 467 ++ .../mir-importer/src/translator/cute/emit.rs | 1323 +++++ .../src/translator/cute/epilogue.rs | 152 + .../src/translator/cute/epilogue_emit.rs | 478 ++ .../src/translator/cute/layout.rs | 152 + .../mir-importer/src/translator/cute/mod.rs | 2070 ++++++++ .../src/translator/cute/pipeline.rs | 206 + .../src/translator/cute/pipeline_emit.rs | 895 ++++ .../src/translator/cute/scheduler.rs | 121 + .../src/translator/cute/scheduler_emit.rs | 540 ++ .../src/translator/cute/smem_mma.rs | 212 + .../src/translator/cute/smem_mma_emit.rs | 1183 +++++ .../src/translator/cute/static_config.rs | 508 ++ .../src/translator/cute/tensor.rs | 243 + .../src/translator/cute/tensor_emit.rs | 710 +++ crates/mir-importer/src/translator/mod.rs | 4 + .../src/translator/rvalue/fn_ptr.rs | 14 + .../src/translator/rvalue/place_read.rs | 5 + .../src/translator/terminator/mod.rs | 24 +- crates/mir-importer/src/translator/types.rs | 23 + crates/mir-importer/src/translator/values.rs | 21 + crates/rustc-codegen-cuda/Cargo.lock | 21 + crates/rustc-codegen-cuda/src/collector.rs | 431 +- cuteir/README.md | 89 + cuteir/cute-layout/Cargo.toml | 10 + cuteir/cute-layout/src/algebra.rs | 1248 +++++ cuteir/cute-layout/src/composed.rs | 534 ++ cuteir/cute-layout/src/lib.rs | 877 ++++ cuteir/cute-layout/src/predication.rs | 479 ++ cuteir/cute-layout/src/swizzle.rs | 386 ++ cuteir/cute-layout/src/validation.rs | 1668 ++++++ cuteir/cute-rs/Cargo.toml | 20 + cuteir/cute-rs/src/block_scaled.rs | 860 ++++ cuteir/cute-rs/src/block_scaled_mma.rs | 965 ++++ cuteir/cute-rs/src/cooperative.rs | 148 + cuteir/cute-rs/src/epilogue.rs | 375 ++ cuteir/cute-rs/src/lib.rs | 99 + cuteir/cute-rs/src/markers.rs | 217 + cuteir/cute-rs/src/mma.rs | 104 + cuteir/cute-rs/src/numeric.rs | 240 + cuteir/cute-rs/src/pipeline.rs | 555 ++ cuteir/cute-rs/src/scheduler.rs | 413 ++ cuteir/cute-rs/src/tensor.rs | 673 +++ cuteir/cute-rs/src/tile.rs | 104 + cuteir/cute-rs/src/tiled_copy.rs | 225 + cuteir/cute-rs/src/tma.rs | 528 ++ .../tests/compile_fail/immutable_store.rs | 15 + .../tests/compile_fail/immutable_store.stderr | 11 + .../tests/compile_fail/invalid_storage.rs | 14 + .../tests/compile_fail/invalid_storage.stderr | 35 + .../compile_fail/mma_operand_role_mismatch.rs | 16 + .../mma_operand_role_mismatch.stderr | 15 + .../compile_fail/mma_scale_role_mismatch.rs | 24 + .../mma_scale_role_mismatch.stderr | 27 + .../tests/compile_fail/tensor_mut_clone.rs | 12 + .../compile_fail/tensor_mut_clone.stderr | 11 + .../tests/compile_fail/tensor_mut_copy.rs | 12 + .../tests/compile_fail/tensor_mut_copy.stderr | 11 + .../tiled_copy_encoding_mismatch.rs | 21 + .../tiled_copy_encoding_mismatch.stderr | 15 + .../compile_fail/tiled_copy_role_mismatch.rs | 23 + .../tiled_copy_role_mismatch.stderr | 15 + .../compile_fail/tma_store_layout_mismatch.rs | 22 + .../tma_store_layout_mismatch.stderr | 15 + cuteir/cute-rs/tests/type_safety.rs | 18 + cuteir/dialect-cute/Cargo.toml | 20 + cuteir/dialect-cute/src/attributes.rs | 1088 ++++ cuteir/dialect-cute/src/epilogue_ops.rs | 1182 +++++ cuteir/dialect-cute/src/gemm_tma_ops.rs | 783 +++ cuteir/dialect-cute/src/gemv_ops.rs | 1095 ++++ cuteir/dialect-cute/src/lib.rs | 142 + cuteir/dialect-cute/src/ops.rs | 1509 ++++++ cuteir/dialect-cute/src/pipeline_ops.rs | 1376 +++++ cuteir/dialect-cute/src/scheduler_ops.rs | 598 +++ cuteir/dialect-cute/src/smem_mma_ops.rs | 1571 ++++++ cuteir/dialect-cute/src/tensor_ops.rs | 952 ++++ cuteir/dialect-cute/src/types.rs | 1222 +++++ cuteir/dialect-cute/src/verify.rs | 4537 +++++++++++++++++ .../examples/blockscale_gemm_cute/Cargo.lock | 567 ++ .../examples/blockscale_gemm_cute/Cargo.toml | 16 + .../examples/blockscale_gemm_cute/README.md | 122 + .../examples/blockscale_gemm_cute/src/main.rs | 684 +++ cuteir/examples/elementwise_cute/Cargo.lock | 566 ++ cuteir/examples/elementwise_cute/Cargo.toml | 16 + cuteir/examples/elementwise_cute/README.md | 93 + cuteir/examples/elementwise_cute/src/main.rs | 167 + cuteir/examples/nvfp4_gemv_cute/Cargo.lock | 566 ++ cuteir/examples/nvfp4_gemv_cute/Cargo.toml | 17 + cuteir/examples/nvfp4_gemv_cute/README.md | 103 + cuteir/examples/nvfp4_gemv_cute/src/main.rs | 423 ++ dependency-licenses.csv | 3 + scripts/check-dependency-licenses.sh | 20 +- scripts/check-example-license-policy.sh | 42 +- 105 files changed, 41400 insertions(+), 45 deletions(-) create mode 100644 crates/mir-importer/src/translator/cute/block_scaled.rs create mode 100644 crates/mir-importer/src/translator/cute/block_scaled_emit.rs create mode 100644 crates/mir-importer/src/translator/cute/emit.rs create mode 100644 crates/mir-importer/src/translator/cute/epilogue.rs create mode 100644 crates/mir-importer/src/translator/cute/epilogue_emit.rs create mode 100644 crates/mir-importer/src/translator/cute/layout.rs create mode 100644 crates/mir-importer/src/translator/cute/mod.rs create mode 100644 crates/mir-importer/src/translator/cute/pipeline.rs create mode 100644 crates/mir-importer/src/translator/cute/pipeline_emit.rs create mode 100644 crates/mir-importer/src/translator/cute/scheduler.rs create mode 100644 crates/mir-importer/src/translator/cute/scheduler_emit.rs create mode 100644 crates/mir-importer/src/translator/cute/smem_mma.rs create mode 100644 crates/mir-importer/src/translator/cute/smem_mma_emit.rs create mode 100644 crates/mir-importer/src/translator/cute/static_config.rs create mode 100644 crates/mir-importer/src/translator/cute/tensor.rs create mode 100644 crates/mir-importer/src/translator/cute/tensor_emit.rs create mode 100644 cuteir/README.md create mode 100644 cuteir/cute-layout/Cargo.toml create mode 100644 cuteir/cute-layout/src/algebra.rs create mode 100644 cuteir/cute-layout/src/composed.rs create mode 100644 cuteir/cute-layout/src/lib.rs create mode 100644 cuteir/cute-layout/src/predication.rs create mode 100644 cuteir/cute-layout/src/swizzle.rs create mode 100644 cuteir/cute-layout/src/validation.rs create mode 100644 cuteir/cute-rs/Cargo.toml create mode 100644 cuteir/cute-rs/src/block_scaled.rs create mode 100644 cuteir/cute-rs/src/block_scaled_mma.rs create mode 100644 cuteir/cute-rs/src/cooperative.rs create mode 100644 cuteir/cute-rs/src/epilogue.rs create mode 100644 cuteir/cute-rs/src/lib.rs create mode 100644 cuteir/cute-rs/src/markers.rs create mode 100644 cuteir/cute-rs/src/mma.rs create mode 100644 cuteir/cute-rs/src/numeric.rs create mode 100644 cuteir/cute-rs/src/pipeline.rs create mode 100644 cuteir/cute-rs/src/scheduler.rs create mode 100644 cuteir/cute-rs/src/tensor.rs create mode 100644 cuteir/cute-rs/src/tile.rs create mode 100644 cuteir/cute-rs/src/tiled_copy.rs create mode 100644 cuteir/cute-rs/src/tma.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/immutable_store.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/immutable_store.stderr create mode 100644 cuteir/cute-rs/tests/compile_fail/invalid_storage.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/invalid_storage.stderr create mode 100644 cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.stderr create mode 100644 cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.stderr create mode 100644 cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.stderr create mode 100644 cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.stderr create mode 100644 cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.stderr create mode 100644 cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr create mode 100644 cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.stderr create mode 100644 cuteir/cute-rs/tests/type_safety.rs create mode 100644 cuteir/dialect-cute/Cargo.toml create mode 100644 cuteir/dialect-cute/src/attributes.rs create mode 100644 cuteir/dialect-cute/src/epilogue_ops.rs create mode 100644 cuteir/dialect-cute/src/gemm_tma_ops.rs create mode 100644 cuteir/dialect-cute/src/gemv_ops.rs create mode 100644 cuteir/dialect-cute/src/lib.rs create mode 100644 cuteir/dialect-cute/src/ops.rs create mode 100644 cuteir/dialect-cute/src/pipeline_ops.rs create mode 100644 cuteir/dialect-cute/src/scheduler_ops.rs create mode 100644 cuteir/dialect-cute/src/smem_mma_ops.rs create mode 100644 cuteir/dialect-cute/src/tensor_ops.rs create mode 100644 cuteir/dialect-cute/src/types.rs create mode 100644 cuteir/dialect-cute/src/verify.rs create mode 100644 cuteir/examples/blockscale_gemm_cute/Cargo.lock create mode 100644 cuteir/examples/blockscale_gemm_cute/Cargo.toml create mode 100644 cuteir/examples/blockscale_gemm_cute/README.md create mode 100644 cuteir/examples/blockscale_gemm_cute/src/main.rs create mode 100644 cuteir/examples/elementwise_cute/Cargo.lock create mode 100644 cuteir/examples/elementwise_cute/Cargo.toml create mode 100644 cuteir/examples/elementwise_cute/README.md create mode 100644 cuteir/examples/elementwise_cute/src/main.rs create mode 100644 cuteir/examples/nvfp4_gemv_cute/Cargo.lock create mode 100644 cuteir/examples/nvfp4_gemv_cute/Cargo.toml create mode 100644 cuteir/examples/nvfp4_gemv_cute/README.md create mode 100644 cuteir/examples/nvfp4_gemv_cute/src/main.rs diff --git a/Cargo.lock b/Cargo.lock index 5f5d537bbb..f9ff24019f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -484,6 +484,7 @@ name = "cuda-oxide-codegen" version = "0.1.0" dependencies = [ "cuda-target-spec", + "dialect-cute", "dialect-iket", "dialect-mir", "dialect-nvvm", @@ -505,6 +506,19 @@ 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 = [ + "cuda-core", + "cuda-device", + "cute-layout", + "trybuild", +] + name = "dashmap" version = "6.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -518,6 +532,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", +] + [[package]] name = "dialect-iket" version = "0.1.0" @@ -992,6 +1020,8 @@ version = "0.2.1" dependencies = [ "combine", "cuda-oxide-codegen", + "cute-layout", + "dialect-cute", "dialect-iket", "dialect-mir", "dialect-nvvm", diff --git a/Cargo.toml b/Cargo.toml index 2c8892aeb0..01ba4f4347 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,10 @@ [workspace] resolver = "3" members = [ + # CuTe frontend and semantic IR shared by every CuTe backend. + "cuteir/cute-layout", + "cuteir/cute-rs", + "cuteir/dialect-cute", # Core crates "crates/cuda-target-spec", "crates/cuda-intrinsics", @@ -45,7 +49,7 @@ repository = "https://github.com/NVlabs/cuda-oxide.git" # Note: crates/rustc-codegen-cuda is NOT a workspace member because it requires # special rustc nightly features and a different 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 cuteir/examples/ [workspace.dependencies] # Pliron MLIR framework, pinned to upstream struct-layout, GEP no-wrap, and @@ -63,6 +67,8 @@ pliron-llvm = { git = "https://github.com/pliron-org/pliron", rev = "edd41fe0615 # Internal crates cuda-target-spec = { path = "crates/cuda-target-spec" } llvm-export = { path = "crates/llvm-export" } +cute-layout = { path = "cuteir/cute-layout" } +dialect-cute = { path = "cuteir/dialect-cute" } dialect-mir = { path = "crates/dialect-mir" } dialect-iket = { path = "crates/dialect-iket" } dialect-ptx = { path = "crates/dialect-ptx" } diff --git a/README.md b/README.md index 2b1f424f2b..ec171a87e0 100644 --- a/README.md +++ b/README.md @@ -297,6 +297,8 @@ cargo oxide run gemm_sol_final | `cuda-bindings` | Raw `bindgen` FFI bindings to `cuda.h` (shared with cutile-rs) | | `cuda-core` | Safe RAII wrappers (`CudaContext`, `DeviceBuffer`, ...); SIMT API under `cuda_core::simt` | | `cuda-async` | Async layer (`DeviceOperation`, `DeviceBox`, ...); SIMT API under `cuda_async::simt` | +| `cute-layout` | CuTe-style layout algebra and coordinate mapping | +| `cute-rs` | Typed CuTe tensors, tiles, copies, pipelines, and MMA operations | | `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`) | @@ -310,6 +312,7 @@ cargo oxide run gemm_sol_final | `mir-importer` | Rust MIR -> `dialect-mir` translation + pipeline | | `mir-lower` | `dialect-mir` -> LLVM dialect lowering | | `dialect-mir` | pliron dialect modelling Rust MIR | +| `dialect-cute` | Backend-neutral semantic IR for CuTe tensor programs | | `dialect-iket` | pliron dialect modelling in-kernel event tracing | | `iket-lower` | `dialect-iket` profiles + instrumentation lowering | | `llvm-export` | pliron-llvm shim + textual `.ll` exporter | diff --git a/crates/cargo-oxide/src/commands/context.rs b/crates/cargo-oxide/src/commands/context.rs index 93b62ecb99..3b8f9541b0 100644 --- a/crates/cargo-oxide/src/commands/context.rs +++ b/crates/cargo-oxide/src/commands/context.rs @@ -490,24 +490,53 @@ fn parse_config_env( /// Resolve an example name to its directory path, or exit with a list of /// available examples if not found. +pub(super) fn example_roots(ctx: &Context) -> Vec { + let mut roots = vec![ctx.examples_dir.clone()]; + if ctx.is_workspace { + let cute_examples = ctx.workspace_root.join("cuteir/examples"); + if cute_examples.is_dir() { + roots.push(cute_examples); + } + } + roots +} + 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); - } + 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/crates/cargo-oxide/src/commands/examples_list.rs b/crates/cargo-oxide/src/commands/examples_list.rs index 7eff89cb16..a7e5a97d48 100644 --- a/crates/cargo-oxide/src/commands/examples_list.rs +++ b/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,11 @@ 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 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 +53,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/crates/cargo-oxide/src/commands/tests.rs b/crates/cargo-oxide/src/commands/tests.rs index da2f29f6a9..bd4e9dcb95 100644 --- a/crates/cargo-oxide/src/commands/tests.rs +++ b/crates/cargo-oxide/src/commands/tests.rs @@ -3951,6 +3951,49 @@ 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_include_existing_cute_examples() { + let root = unique_temp_dir("cargo_oxide_cute_example_root"); + let standard = root.join("crates/rustc-codegen-cuda/examples"); + let cute = root.join("cuteir/examples"); + std::fs::create_dir_all(&standard).unwrap(); + std::fs::create_dir_all(&cute).unwrap(); + + let mut ctx = test_context(OxideConfig::default()); + ctx.workspace_root = root.clone(); + ctx.examples_dir = standard.clone(); + ctx.is_workspace = true; + assert_eq!(example_roots(&ctx), [standard, cute]); + + 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"); diff --git a/crates/cargo-oxide/src/main.rs b/crates/cargo-oxide/src/main.rs index 8572e66eaf..16a6136aee 100644 --- a/crates/cargo-oxide/src/main.rs +++ b/crates/cargo-oxide/src/main.rs @@ -1117,7 +1117,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/ and cuteir/examples/" + ); std::process::exit(1); } diff --git a/crates/cuda-oxide-codegen/Cargo.toml b/crates/cuda-oxide-codegen/Cargo.toml index 0b2b091e6e..8670166103 100644 --- a/crates/cuda-oxide-codegen/Cargo.toml +++ b/crates/cuda-oxide-codegen/Cargo.toml @@ -18,6 +18,7 @@ thiserror = { workspace = true } dialect-mir = { workspace = true } dialect-iket = { workspace = true } dialect-nvvm = { workspace = true } +dialect-cute = { workspace = true } iket-lower = { workspace = true } mir-lower = { workspace = true } mir-transforms = { workspace = true } diff --git a/crates/cuda-oxide-codegen/src/pipeline.rs b/crates/cuda-oxide-codegen/src/pipeline.rs index 0fbfabdc31..649ea57db7 100644 --- a/crates/cuda-oxide-codegen/src/pipeline.rs +++ b/crates/cuda-oxide-codegen/src/pipeline.rs @@ -259,6 +259,15 @@ pub fn compile_translated_module( .emit(format!("{}", module.deref(ctx).disp(ctx))); } + // This is the shared CuTe backend seam. Validate the complete semantic + // story once, before either a native lowering or an external compiler + // translation is allowed to consume it. + dialect_cute::verify::verify_cute_semantics(ctx, module).map_err(|error| { + PipelineError::Lowering(format!( + "dialect-cute semantic verification failed: {error}" + )) + })?; + // 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). diff --git a/crates/cuda-oxide-codegen/src/prep.rs b/crates/cuda-oxide-codegen/src/prep.rs index 86a8302a07..998559ab87 100644 --- a/crates/cuda-oxide-codegen/src/prep.rs +++ b/crates/cuda-oxide-codegen/src/prep.rs @@ -6,9 +6,48 @@ 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::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 +75,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 +88,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 +194,32 @@ pub fn prepare_mir_module( "module post-borrowed-pointer-read-canonicalization", )?; + if has_high_level_cute_operations { + // A reborrow of a local CuTe carrier binding (`&mut tensor`) imports + // as a pointer-kind-only `mir.cast` on the alloca slot pointer. That + // cast is a non-promotable use for mem2reg, but ghost CuTe values are + // forbidden from living in memory at all, so the kind refinement + // carries no ABI meaning for them. Once the first mem2reg pass has + // dissolved the pointer slots, the cast's remaining uses are plain + // memory addresses; fold it back onto the slot pointer so the second + // pass can dissolve the ghost slot itself. + 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 +238,149 @@ 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) + { + 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) { + if 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 +424,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 +504,314 @@ 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 + // `store_linear(&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 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/crates/mir-importer/Cargo.toml b/crates/mir-importer/Cargo.toml index 3f32e242a4..ff7cc7d5b1 100644 --- a/crates/mir-importer/Cargo.toml +++ b/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/crates/mir-importer/src/translator/cute/block_scaled.rs b/crates/mir-importer/src/translator/cute/block_scaled.rs new file mode 100644 index 0000000000..74734b7c5c --- /dev/null +++ b/crates/mir-importer/src/translator/cute/block_scaled.rs @@ -0,0 +1,254 @@ +/* + * 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. + +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}; + +const BLOCK_SCALED_TENSOR: &str = "cute_rs::block_scaled::BlockScaledTensor"; +const BLOCK_SCALED_THREAD_ROW: &str = "cute_rs::block_scaled::BlockScaledThreadRow"; +const BLOCK_SCALED_TILE_64: &str = "cute_rs::block_scaled::BlockScaledTile64"; +const LOADED_BLOCK_SCALED_TILE_64: &str = "cute_rs::block_scaled::LoadedBlockScaledTile64"; +const K_MAJOR: &str = "cute_rs::block_scaled::KMajor"; +const MKL: &str = "cute_rs::block_scaled::Mkl"; +const NKL: &str = "cute_rs::block_scaled::Nkl"; +const E2M1: &str = "cute_rs::numeric::E2M1"; +const UE8M0: &str = "cute_rs::numeric::UE8M0"; + +/// Which stage of the block-scaled view chain one Rust ADT represents. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum BlockScaledStage { + Full, + Row, + KTile64, + Fragment64, +} + +/// Static facts used to classify calls and build ghost types. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct BlockScaledRust { + pub role: CuteTensorRoleAttr, + pub stage: BlockScaledStage, +} + +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)) +} + +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) +} + +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}`" + )), + } +} + +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)?, E2M1)?; + exact_marker(&type_arg(&args, 2, BLOCK_SCALED_TENSOR)?, UE8M0)?; + 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)?)?, + 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)?)?, + stage: if path == BLOCK_SCALED_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)?)?, + 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), + } +} + +fn scaled_view_type( + ctx: &mut Context, + role: CuteTensorRoleAttr, + 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, + CuteTensorFormatAttr::UE8M0, + 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, CuteScaledLayoutAttr::Full), + BlockScaledStage::Row => scaled_view_type(ctx, decoded.role, CuteScaledLayoutAttr::Row), + BlockScaledStage::KTile64 => { + scaled_view_type(ctx, decoded.role, CuteScaledLayoutAttr::KTile(64)) + } + BlockScaledStage::Fragment64 => { + let source = scaled_view_type(ctx, decoded.role, CuteScaledLayoutAttr::KTile(64)); + CuteFragmentType::get(ctx, source).into() + } + })) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn stage_names_keep_the_source_flow_visible() { + assert_ne!(BlockScaledStage::Full, BlockScaledStage::Row); + assert_ne!(BlockScaledStage::KTile64, BlockScaledStage::Fragment64); + } + + #[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, + 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)); + } +} diff --git a/crates/mir-importer/src/translator/cute/block_scaled_emit.rs b/crates/mir-importer/src/translator/cute/block_scaled_emit.rs new file mode 100644 index 0000000000..e3c222f231 --- /dev/null +++ b/crates/mir-importer/src/translator/cute/block_scaled_emit.rs @@ -0,0 +1,467 @@ +/* + * 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::CuteTensorRoleAttr; +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::{Located, 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::block_scaled::{BlockScaledStage, decode_block_scaled}; + +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); + } +} + +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)) +} + +#[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) +} + +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) +} + +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)) +} + +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.role) +} + +#[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 role = destination_facts( + body, + destination, + BlockScaledStage::Full, + &loc, + "block_scaled_make", + )?; + + 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 = 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", + ) +} + +#[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", + ) +} + +#[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", + ) +} + +#[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", + ) +} + +#[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/crates/mir-importer/src/translator/cute/emit.rs b/crates/mir-importer/src/translator/cute/emit.rs new file mode 100644 index 0000000000..2602c2eec7 --- /dev/null +++ b/crates/mir-importer/src/translator/cute/emit.rs @@ -0,0 +1,1323 @@ +/* + * 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_div` 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::{Located, 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::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)) +} + +#[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; + + 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) + } + + fn pointer_address_space(ctx: &Context, value: Value) -> u32 { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .unwrap() + .address_space + } + + 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) + } + + #[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}"); + } + + #[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::tile::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::tile::assume_div::(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_div 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_div 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_div call without target not supported".to_string()) + ); + }; + Ok(emit_goto(ctx, *target_idx, goto_prev, block_map, loc)) +} + +/// `cute_rs::tile::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::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::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)) +} diff --git a/crates/mir-importer/src/translator/cute/epilogue.rs b/crates/mir-importer/src/translator/cute/epilogue.rs new file mode 100644 index 0000000000..75d9060764 --- /dev/null +++ b/crates/mir-importer/src/translator/cute/epilogue.rs @@ -0,0 +1,152 @@ +/* + * 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. + +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}; + +const EPILOGUE_TILE: &str = "cute_rs::epilogue::Sm120Epilogue128x128"; +const EPILOGUE_WARP_SLICE: &str = "cute_rs::epilogue::Sm120EpilogueWarp128x128"; +const ACCUMULATOR: &str = "cute_rs::block_scaled_mma::Mxf4AccumulatorTile2x8"; +const TMA_STORE_PIPELINE: &str = "cute_rs::pipeline::TmaStorePipeline"; + +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)) +} + +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), + } +} + +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) +} + +pub(crate) fn is_epilogue_tile(ty: &Ty) -> Result { + is_exact_marker(ty, EPILOGUE_TILE) +} + +pub(crate) fn is_epilogue_tile_receiver(ty: &Ty) -> Result { + behind_one_pointer(ty, is_epilogue_tile) +} + +pub(crate) fn is_epilogue_warp_slice(ty: &Ty) -> Result { + is_exact_marker(ty, EPILOGUE_WARP_SLICE) +} + +pub(crate) fn is_epilogue_warp_slice_receiver(ty: &Ty) -> Result { + behind_one_pointer(ty, is_epilogue_warp_slice) +} + +pub(crate) fn is_accumulator(ty: &Ty) -> Result { + behind_one_pointer(ty, |ty| is_exact_marker(ty, ACCUMULATOR)) +} + +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) +} + +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() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[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/crates/mir-importer/src/translator/cute/epilogue_emit.rs b/crates/mir-importer/src/translator/cute/epilogue_emit.rs new file mode 100644 index 0000000000..45f6257094 --- /dev/null +++ b/crates/mir-importer/src/translator/cute/epilogue_emit.rs @@ -0,0 +1,478 @@ +/* + * 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::{Located, Location}; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::r#type::TypeHandle; +use pliron::value::Value; +use rustc_public::mir; + +use super::emit::recover_cta_shared_pointer; +use super::epilogue::{ + is_accumulator, is_epilogue_tile, is_epilogue_tile_receiver, is_epilogue_warp_slice_receiver, + one_unsigned_const, tile_type, +}; +use 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, +}; + +fn invalid(what: &str, error: impl core::fmt::Display) -> pliron::result::Error { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid {what}: {error}" + ))) +} + +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)), + } +} + +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)) +} + +#[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) +} + +#[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)) +} + +#[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", + ) +} + +#[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", + ) +} + +#[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") +} + +#[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") +} + +#[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", + ) +} + +#[derive(Clone, Copy)] +pub(crate) enum StorePipelineEffect { + Acquire, + Commit, + Tail, +} + +#[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/crates/mir-importer/src/translator/cute/layout.rs b/crates/mir-importer/src/translator/cute/layout.rs new file mode 100644 index 0000000000..b808c9819d --- /dev/null +++ b/crates/mir-importer/src/translator/cute/layout.rs @@ -0,0 +1,152 @@ +/* + * 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. +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_div::`, 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_div 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_div 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_div is missing divisor const generic at substs[0]".to_string() + ) + ); + }; + let divisor = const_u64(divisor, "assume_div divisor", loc)?; + if divisor == 0 { + return input_err!( + loc.clone(), + TranslationErr::unsupported("assume_div divisor must be greater than zero".to_string()) + ); + } + Ok(divisor) +} diff --git a/crates/mir-importer/src/translator/cute/mod.rs b/crates/mir-importer/src/translator/cute/mod.rs new file mode 100644 index 0000000000..496322274f --- /dev/null +++ b/crates/mir-importer/src/translator/cute/mod.rs @@ -0,0 +1,2070 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Recognition of `cute-rs` device-library calls. +//! +//! ALL cute-specific importer code lives in this module; the shared +//! translator contains only the one dispatch line in `translate_call`. +//! Mirrors the `terminator/intrinsics/` pattern. +//! +//! Recognition is STRUCTURAL, not string-based: we rebuild the callee's +//! canonical definition path by walking the DefId parent chain (the +//! cuda-intrinsics mechanism) and compare it against [`CUTE_FNS`]. This is +//! immune to the two failure modes of name matching: +//! +//! - rustc's name printer prefers public re-export paths +//! (`cute_rs::load_tile`) over definition paths +//! (`cute_rs::tile::load_tile`), and re-exports are API surface that +//! moves; +//! - trait-method paths print as `::method`, which +//! prefix checks miss entirely. +//! +//! Contract: a recognized call is NEVER translated body-by-body. Static +//! parameters are read from the monomorphized substs after validating the +//! per-function [`SubstKind`] schema. Explicit source-abstraction modules +//! contain ordinary device-safe Rust and fall through to body translation; +//! LLVM then erases their `#[inline(always)]` boundaries. Every other unknown +//! `cute_rs` call is a hard error (body translation of a recognition stub +//! would scalarize the layout: the raising trap). + +pub(crate) mod block_scaled; +pub(crate) mod block_scaled_emit; +pub(crate) mod emit; +pub(crate) mod epilogue; +pub(crate) mod epilogue_emit; +pub(crate) mod layout; +pub(crate) mod pipeline; +pub(crate) mod pipeline_emit; +pub(crate) mod scheduler; +pub(crate) mod scheduler_emit; +pub(crate) mod smem_mma; +pub(crate) mod smem_mma_emit; +pub(crate) mod static_config; +pub(crate) mod tensor; +pub(crate) mod tensor_emit; + +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 rustc_public::ty::FnDef; + +/// Recognized cute-rs entry points. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum CuteFn { + MakeTensorRead, + MakeTensorWrite, + ZippedDivideRead, + ZippedDivideWrite, + SliceRead, + SliceWrite, + TensorTileIsFull, + TensorTileBase, + TensorLoadTile, + TensorStoreTile, + TensorStoreElementAbs, + BlockScaledMake, + BlockScaledThreadRow, + BlockScaledKTile, + BlockScaledLoadK64, + BlockScaledDotK64, + SchedulerNew1d, + SchedulerHasWork, + SchedulerCurrentTile, + WorkTileCoordinates, + SchedulerAdvance, + TmaLoadPipelineMake, + TmaLoadPipelineFromRawPartsRejected, + TmaLoadPipelineInit, + PipelineStateNew, + PipelineStateSlot, + PipelineStateAdvance, + PipelineProducerAcquire, + PipelineProducerExpectTx, + PipelineConsumerWait, + PipelineConsumerRelease, + PipelineProducerTail, + SharedTensorOverlay, + TiledMmaSlice, + MmaFragmentFill, + MmaLoadScales, + MmaFragmentSliceK, + MmaLoadA, + MmaPartitionB, + TiledGemm, + EpilogueOverlay, + EpilogueWarpSlice, + EpilogueStoreFragment, + EpilogueSyncReusable, + EpilogueSyncReadyForTma, + EpilogueHalf, + TmaStoreAcquire, + TmaStoreCommit, + TmaStoreTail, + LoadTile, + StoreTile, + AssumeDiv, + CopyG2S, + LoadMatrixA, + LoadMatrixB, + CopyTma2d, + 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 { + Type, + Const, + Lifetime, +} + +/// One recognized function: its canonical DEFINITION path (what the DefId +/// parent walk reconstructs; re-exports never appear here) and the shape +/// its generics must have. +pub(crate) struct CuteFnSpec { + pub canonical: &'static str, + pub func: CuteFn, + pub schema: &'static [SubstKind], +} + +/// The recognition table. Grows by one entry per recognized function; when +/// the op set stabilizes this graduates to catalog generation (the +/// cuda-intrinsics-gen model) and these entries become generated. +pub(crate) const CUTE_FNS: &[CuteFnSpec] = &[ + 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::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::tiled_copy::__compiler::shared_tensor_overlay", + func: CuteFn::SharedTensorOverlay, + schema: &[ + SubstKind::Type, + SubstKind::Type, + SubstKind::Type, + SubstKind::Type, + ], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::tiled_mma_slice", + func: CuteFn::TiledMmaSlice, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::fragment_fill", + func: CuteFn::MmaFragmentFill, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::mma_load_scales", + func: CuteFn::MmaLoadScales, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::fragment_slice_k", + func: CuteFn::MmaFragmentSliceK, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::mma_load_a", + func: CuteFn::MmaLoadA, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::mma_partition_b", + func: CuteFn::MmaPartitionB, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::tiled_gemm", + func: CuteFn::TiledGemm, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::__compiler::epilogue_smem_overlay", + func: CuteFn::EpilogueOverlay, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::__compiler::epilogue_warp_slice", + func: CuteFn::EpilogueWarpSlice, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::epilogue_store_fragment", + func: CuteFn::EpilogueStoreFragment, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::__compiler::epilogue_sync_reusable", + func: CuteFn::EpilogueSyncReusable, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::__compiler::epilogue_sync_ready_for_tma", + func: CuteFn::EpilogueSyncReadyForTma, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::__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::tile::load_tile", + func: CuteFn::LoadTile, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tile::store_tile", + func: CuteFn::StoreTile, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tile::assume_div", + 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::tma::copy_tma_2d", + func: CuteFn::CopyTma2d, + schema: &[SubstKind::Type, SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::tma::copy_tma_s2g_2d", + func: CuteFn::CopyTmaS2g2d, + schema: &[SubstKind::Type, SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::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, + ], + }, +]; + +enum CuteIdentity { + /// Callee is not from the cute_rs crate; normal translation proceeds. + NotCuteCrate, + /// A recognized entry point. + Known(&'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(String), +} + +/// Definition modules containing ordinary Rust source abstractions rather +/// than importer-owned recognition stubs. +/// +/// This is deliberately a closed list. Adding a module means every function +/// defined directly or in an impl below it must have a real, device-safe body. +/// Panic-only recognition stubs are forbidden in these modules: they would be +/// body-translated instead of dispatched through [`CUTE_FNS`]. New modules +/// therefore require an explicit audit and an update to the exact-list test. +const CUTE_SOURCE_ABSTRACTION_MODULES: &[&str] = &[ + "block_scaled", + "block_scaled_mma", + "epilogue", + "numeric", + "pipeline", + "scheduler", + "tensor", + "tiled_copy", +]; + +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; + }; + let first_module = rest.split("::").next().unwrap_or_default(); + CUTE_SOURCE_ABSTRACTION_MODULES.contains(&first_module) +} + +/// Rebuild the canonical definition path by walking DefId parents, +/// mirroring `intrinsics::generated::classify_raw_intrinsic`. +fn canonical_path(fn_def: &FnDef, crate_name: &str) -> String { + let mut segments = Vec::new(); + let mut current = Some(fn_def.def_id()); + while let Some(def_id) = current { + let printed = def_id.name(); + let segment = printed.as_str().rsplit("::").next().unwrap_or_default(); + segments.push(segment.to_owned()); + current = def_id.parent(); + } + canonical_path_from_leaf_segments(crate_name, segments) +} + +/// Join a DefId parent walk without erasing a same-named nested module. +pub(super) fn canonical_path_from_leaf_segments( + crate_name: &str, + mut segments: Vec, +) -> String { + segments.reverse(); + // Only the first segment is the crate root. A nested module may legally + // have the same spelling and must remain part of the identity: + // + // cute_rs::cute_rs::tile::load_tile != cute_rs::tile::load_tile + if segments + .first() + .is_some_and(|segment| segment == crate_name) + { + segments.remove(0); + } + format!("{crate_name}::{}", segments.join("::")) +} + +#[cfg(test)] +mod canonical_path_tests { + use super::{ + block_scaled_method_name_from_path, canonical_path_from_leaf_segments, + is_source_abstraction_path, pipeline_method_name_from_path, + scheduler_method_name_from_path, smem_mma_method_name_from_path, + tensor_method_name_from_path, + }; + + #[test] + fn removes_only_the_crate_root_segment() { + let path = canonical_path_from_leaf_segments( + "cute_rs", + ["load_tile", "tile", "cute_rs", "cute_rs"] + .into_iter() + .map(str::to_string) + .collect(), + ); + assert_eq!(path, "cute_rs::cute_rs::tile::load_tile"); + } + + #[test] + fn source_abstraction_modules_are_exact_and_closed() { + assert_eq!( + super::CUTE_SOURCE_ABSTRACTION_MODULES, + &[ + "block_scaled", + "block_scaled_mma", + "epilogue", + "numeric", + "pipeline", + "scheduler", + "tensor", + "tiled_copy", + ] + ); + assert!(is_source_abstraction_path( + "cute_rs::block_scaled::BlockScaledTensor<...>::thread_row", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute-rs::numeric::PackedE2M1x2::to_f32x2", + "cute-rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::tensor::Tensor<...>::layout", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::tiled_copy::TiledCopy<...>::copy", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::block_scaled_mma::Mxfp4TiledMma<...>::gemm", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::epilogue::Sm120EpilogueWarp128x128::store_atom", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::pipeline::TmaLoadPipeline<...>::consumer_wait", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::scheduler::StaticPersistentTileScheduler<...>::current", + "cute_rs" + )); + assert!(!is_source_abstraction_path( + "cute_rs::tile::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" + )); + } + + #[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_linear"), + Some("store_linear") + ); + 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 + ); + } + + #[test] + fn block_scaled_method_path_ignores_only_the_impl_owner_spelling() { + assert_eq!( + block_scaled_method_name_from_path("cute_rs::block_scaled::KMajor>::from_slices"), + Some("from_slices") + ); + assert_eq!( + block_scaled_method_name_from_path("cute_rs::block_scaled::{impl#21}::dot_accumulate"), + Some("dot_accumulate") + ); + assert_eq!( + block_scaled_method_name_from_path( + "cute_rs::block_scaled::__compiler::block_scaled_dot_k64" + ), + Some("block_scaled_dot_k64"), + "the exact helper table takes precedence over method classification" + ); + assert_eq!( + block_scaled_method_name_from_path("cute_rs::other::{impl#1}::load"), + None + ); + } + + #[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}::coordinates"), + Some("coordinates") + ); + 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" + ); + } + + #[test] + fn pipeline_method_path_ignores_only_the_impl_owner_spelling() { + assert_eq!( + pipeline_method_name_from_path( + "cute_rs::pipeline::TmaLoadPipeline<...>>::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" + ); + } + + #[test] + fn shared_mma_method_path_keeps_module_and_method_exact() { + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::block_scaled_mma::Mxfp4TiledMma<...>>::load_a_128", + "block_scaled_mma" + ), + Some("load_a_128") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute-rs::tiled_copy::{impl#7}::from_raw_parts", + "tiled_copy" + ), + Some("from_raw_parts") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::other::Mxfp4TiledMma<...>>::load_a_128", + "block_scaled_mma" + ), + None + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::block_scaled_mma::load_a_128", + "block_scaled_mma" + ), + None, + "a free function is not an inherent method" + ); + } + + #[test] + fn epilogue_method_paths_keep_each_owner_module_exact() { + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::epilogue::Sm120Epilogue128x128::sync_ready_for_tma", + "epilogue" + ), + Some("sync_ready_for_tma") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::block_scaled_mma::{impl#9}::store_tile", + "block_scaled_mma" + ), + Some("store_tile") + ); + assert_eq!( + pipeline_method_name_from_path( + "cute_rs::pipeline::TmaStorePipeline<1>::producer_commit" + ), + Some("producer_commit") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::epilogue::Sm120EpilogueWarp128x128::store_tile", + "block_scaled_mma" + ), + None, + "store_tile is defined in block_scaled_mma, not epilogue" + ); + } +} + +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), + } +} + +/// Recognize only the ergonomic inherent methods that expose the v0 tensor +/// flow. The DefId's direct parent is an implementation block, whose printed +/// number is intentionally ignored; the receiver/result ADT supplies the +/// stable identity instead. +fn tensor_method_name_from_path(path: &str) -> Option<&str> { + let mut segments = path.split("::"); + let crate_name = segments.next()?; + let module = segments.next()?; + // rustc_public currently prints an impl parent using the tail of its self + // type (for example `Tile1D>`), while older revisions exposed an + // `{impl#N}` segment. That spelling is not identity. The exact receiver + // and result ADTs checked by `classify_tensor_method_call` are. + let impl_owner = segments.next()?; + let method = segments.next()?; + if segments.next().is_some() + || !matches!(crate_name, "cute_rs" | "cute-rs") + || module != "tensor" + || impl_owner.is_empty() + { + return None; + } + Some(method) +} + +/// Return an inherent method name from the exact block-scaled module. +/// +/// rustc's spelling for the implementation owner is not stable. Receiver and +/// result ADTs below provide the identity instead. +fn block_scaled_method_name_from_path(path: &str) -> Option<&str> { + let mut segments = path.split("::"); + let crate_name = segments.next()?; + let module = segments.next()?; + let impl_owner = segments.next()?; + let method = segments.next()?; + if segments.next().is_some() + || !matches!(crate_name, "cute_rs" | "cute-rs") + || module != "block_scaled" + || impl_owner.is_empty() + { + return None; + } + Some(method) +} + +/// Return an inherent method name from the exact shared-copy or MMA module. +fn smem_mma_method_name_from_path<'a>(path: &'a str, expected_module: &str) -> Option<&'a str> { + let mut segments = path.split("::"); + let crate_name = segments.next()?; + let module = segments.next()?; + let impl_owner = segments.next()?; + let method = segments.next()?; + if segments.next().is_some() + || !matches!(crate_name, "cute_rs" | "cute-rs") + || module != expected_module + || impl_owner.is_empty() + { + return None; + } + Some(method) +} + +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, "tiled_copy") == 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, "block_scaled_mma")?; + 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) + } + "accumulate_k64" + if args.len() == 5 + && operand_ty(0) + .is_some_and(|ty| smem_mma::is_accumulator(&ty).unwrap_or(false)) => + { + Some(CuteFn::TiledGemm) + } + _ => None, + } +} + +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") { + 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, "block_scaled_mma") == 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. +fn scheduler_method_name_from_path(path: &str) -> Option<&str> { + let mut segments = path.split("::"); + let crate_name = segments.next()?; + let module = segments.next()?; + let impl_owner = segments.next()?; + let method = segments.next()?; + if segments.next().is_some() + || !matches!(crate_name, "cute_rs" | "cute-rs") + || module != "scheduler" + || impl_owner.is_empty() + { + return None; + } + Some(method) +} + +/// 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. +fn pipeline_method_name_from_path(path: &str) -> Option<&str> { + let mut segments = path.split("::"); + let crate_name = segments.next()?; + let module = segments.next()?; + let impl_owner = segments.next()?; + let method = segments.next()?; + if segments.next().is_some() + || !matches!(crate_name, "cute_rs" | "cute-rs") + || module != "pipeline" + || impl_owner.is_empty() + { + return None; + } + Some(method) +} + +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), + "slot" 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, + } +} + +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) + } + "current_tile" + if args.len() == 1 + && matches!((receiver_scheduler, destination_tile), (Some(source), Some(result)) if source == result) => + { + Some(CuteFn::SchedulerCurrentTile) + } + "coordinates" if args.len() == 1 && receiver_tile.is_some() => { + Some(CuteFn::WorkTileCoordinates) + } + "advance" if args.len() == 1 && receiver_scheduler.is_some() => { + Some(CuteFn::SchedulerAdvance) + } + _ => None, + } +} + +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, + } +} + +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, + }, + "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) + } + "store_linear" + if receiver_view.is_some_and(|view| { + view.access == CuteTensorAccessAttr::ReadWrite + && matches!(view.layout, CuteTensorLayoutAttr::Tile1D(_)) + }) => + { + Some(CuteFn::TensorStoreElementAbs) + } + _ => None, + } +} + +/// 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` when the callee is not a cute-rs function or is an explicit +/// source abstraction (the caller falls through to normal translation). +/// Returns `Some(result)` when this module owns the call, in which case the +/// shared path must not touch it. +#[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 be cute calls (indirect dispatch of cute_rs items is rejected by + // `requires_direct_dispatch` in the function-item path). + 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::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::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( + "TmaLoadPipeline::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::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::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(), + ), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[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/crates/mir-importer/src/translator/cute/pipeline.rs b/crates/mir-importer/src/translator/cute/pipeline.rs new file mode 100644 index 0000000000..9bfac9ada6 --- /dev/null +++ b/crates/mir-importer/src/translator/cute/pipeline.rs @@ -0,0 +1,206 @@ +/* + * 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 +//! TmaLoadPipeline -> !cute.tma_load_pipeline +//! PipelineState -> { slot: u32, phase: u32 } +//! ``` + +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}; + +const TMA_LOAD_PIPELINE: &str = "cute_rs::pipeline::TmaLoadPipeline"; +const PIPELINE_STATE: &str = "cute_rs::pipeline::PipelineState"; +const PRODUCER: &str = "cute_rs::pipeline::Producer"; +const CONSUMER: &str = "cute_rs::pipeline::Consumer"; + +/// Static facts carried by `TmaLoadPipeline`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct TmaLoadPipelineRust { + pub stages: u64, + pub consumer_warps: u32, + pub transaction_bytes: u32, +} + +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(TmaLoadPipelineRust { + stages, + consumer_warps, + transaction_bytes, + })) +} + +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(()) +} + +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), + } +} + +fn pipeline_type(ctx: &mut Context, facts: TmaLoadPipelineRust) -> 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}") + )))), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn load_pipeline_type_keeps_all_static_facts() { + let mut ctx = Context::new(); + dialect_cute::register(&mut ctx); + let facts = TmaLoadPipelineRust { + 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)); + } + + #[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/crates/mir-importer/src/translator/cute/pipeline_emit.rs b/crates/mir-importer/src/translator/cute/pipeline_emit.rs new file mode 100644 index 0000000000..3d227828ec --- /dev/null +++ b/crates/mir-importer/src/translator/cute/pipeline_emit.rs @@ -0,0 +1,895 @@ +/* + * 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 -> slot -> 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::{Located, 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::pipeline::{ + TmaLoadPipelineRust, decode_pipeline_state, decode_pipeline_state_receiver, + decode_tma_load_pipeline, +}; + +const BARRIER_BYTES: u64 = 8; + +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); + } +} + +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)) +} + +#[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) +} + +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) + ))) + }) +} + +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) +} + +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 TmaLoadPipeline receiver" + ))) + }) +} + +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" + ))) + }) +} + +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(()) +} + +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))) +} + +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)) +} + +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) +} + +struct StateValues { + slot: Value, + phase: Value, + slot_pointer: Option, + phase_pointer: Option, + anchor: Ptr, +} + +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"), + }) +} + +#[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 TmaLoadPipeline".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") +} + +#[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") +} + +#[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", + ) +} + +#[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", + ) +} + +#[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", + ) +} + +enum Lifecycle { + ProducerAcquire, + ProducerExpectTx, + ConsumerWait, + ConsumerRelease, + ProducerTail, +} + +#[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) + } +} + +macro_rules! lifecycle_emitter { + ($name:ident, $kind:expr) => { + #[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); + +#[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/crates/mir-importer/src/translator/cute/scheduler.rs b/crates/mir-importer/src/translator/cute/scheduler.rs new file mode 100644 index 0000000000..84017bde1c --- /dev/null +++ b/crates/mir-importer/src/translator/cute/scheduler.rs @@ -0,0 +1,121 @@ +/* + * 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 +//! │ current_tile +//! ▼ +//! WorkTile !cute.work_tile> +//! ``` + +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}; + +const SCHEDULER: &str = "cute_rs::scheduler::StaticPersistentTileScheduler"; +const WORK_TILE: &str = "cute_rs::scheduler::WorkTile"; + +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)) +} + +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() +} + +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), + } +} + +pub(crate) fn decode_scheduler_receiver(ty: &Ty) -> Result, String> { + decode_behind_one_pointer(ty, decode_scheduler) +} + +pub(crate) fn decode_work_tile_receiver(ty: &Ty) -> Result, String> { + decode_behind_one_pointer(ty, decode_work_tile) +} + +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}") + )))), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[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/crates/mir-importer/src/translator/cute/scheduler_emit.rs b/crates/mir-importer/src/translator/cute/scheduler_emit.rs new file mode 100644 index 0000000000..e1cd3d1e44 --- /dev/null +++ b/crates/mir-importer/src/translator/cute/scheduler_emit.rs @@ -0,0 +1,540 @@ +/* + * 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 `WorkTile` is a short-lived compiler handle: +//! +//! ```text +//! new_1d -> has_work -> current_tile -> coordinates -> advance +//! ``` + +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::{Located, 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::scheduler::{ + decode_scheduler, decode_scheduler_receiver, decode_work_tile, decode_work_tile_receiver, +}; + +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); + } +} + +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)) +} + +#[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) +} + +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) + ))) + }) +} + +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) +} + +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" + ))) + }) +} + +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)) +} + +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) +} + +#[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", + ) +} + +#[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", + ) +} + +#[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", + ) +} + +#[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.linear(ctx), + 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", + ) +} + +#[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/crates/mir-importer/src/translator/cute/smem_mma.rs b/crates/mir-importer/src/translator/cute/smem_mma.rs new file mode 100644 index 0000000000..942c167913 --- /dev/null +++ b/crates/mir-importer/src/translator/cute/smem_mma.rs @@ -0,0 +1,212 @@ +/* + * 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. + +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}; + +const SHARED_TENSOR: &str = "cute_rs::tiled_copy::SharedTensor"; +const MXF4_E2M1: &str = "cute_rs::numeric::Mxf4E2M1"; +const UE8M0_X4: &str = "cute_rs::numeric::UE8M0x4"; +const SM120_SCALE_ATOM: &str = "cute_rs::block_scaled::Sm120ScaleAtom"; +const MKL: &str = "cute_rs::block_scaled::Mkl"; +const NKL: &str = "cute_rs::block_scaled::Nkl"; +const TILED_MMA: &str = "cute_rs::block_scaled_mma::Mxfp4TiledMma"; +const ACCUMULATOR: &str = "cute_rs::block_scaled_mma::Mxf4AccumulatorTile2x8"; +const SCALE_STAGE: &str = "cute_rs::block_scaled_mma::Mxf4ScaleTile128"; +const SCALE_K64: &str = "cute_rs::block_scaled_mma::Mxf4ScalePairs128"; +const B_TILE_K64: &str = "cute_rs::block_scaled_mma::Mxf4BTileK64"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum SharedTensorKind { + Data, + Scale, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct SharedTensorRust { + pub kind: SharedTensorKind, + pub role: CuteTensorRoleAttr, + pub placement: ComposedLayout, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct TiledMmaRust { + pub placement: ComposedLayout, +} + +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)) +} + +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) +} + +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}`")) + } +} + +fn is_f16(ty: &Ty) -> bool { + matches!(ty.kind(), TyKind::RigidTy(RigidTy::Float(FloatTy::F16))) +} + +fn is_u32(ty: &Ty) -> bool { + matches!(ty.kind(), TyKind::RigidTy(RigidTy::Uint(UintTy::U32))) +} + +fn scale_placement() -> ComposedLayout { + let inner = "(32,4):(4,1)" + .parse() + .expect("fixed SM120 scale placement is valid"); + ComposedLayout::from_layout(inner, OffsetUnit::Elements) +} + +pub(crate) fn canonical_data_placement() -> ComposedLayout { + let inner: Layout = "(128,32):(32,1)" + .parse() + .expect("fixed SM120 packed-data placement is valid"); + ComposedLayout::new(Swizzle::new(2, 3, 3), 0, inner, OffsetUnit::Elements) + .expect("fixed SM120 packed-data swizzle is valid") +} + +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()) +} + +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), + } +} + +pub(crate) fn decode_shared_tensor_receiver(ty: &Ty) -> Result, String> { + behind_one_pointer(ty, decode_shared_tensor) +} + +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)?)?, + })) +} + +pub(crate) fn decode_tiled_mma_receiver(ty: &Ty) -> Result, String> { + behind_one_pointer(ty, decode_tiled_mma) +} + +fn exact_adt(ty: &Ty, expected: &str) -> Result { + exact_marker(ty, expected) +} + +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)) +} + +pub(crate) fn is_scale_stage(ty: &Ty) -> Result { + exact_adt(ty, SCALE_STAGE) +} + +pub(crate) fn is_scale_k64(ty: &Ty) -> Result { + exact_adt(ty, SCALE_K64) +} + +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)?)?, + })) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn fixed_scale_placement_has_one_128_word_atom() { + assert_eq!(scale_placement().inner().checked_size(), Some(128)); + } + + #[test] + fn canonical_data_placement_locks_the_128_by_32_swizzled_tile() { + let placement = canonical_data_placement(); + assert_eq!(placement.inner().checked_size(), Some(128 * 32)); + assert_eq!(placement.outer(), Swizzle::new(2, 3, 3)); + assert_eq!(placement.offset(), 0); + assert_eq!(placement.unit(), OffsetUnit::Elements); + } +} diff --git a/crates/mir-importer/src/translator/cute/smem_mma_emit.rs b/crates/mir-importer/src/translator/cute/smem_mma_emit.rs new file mode 100644 index 0000000000..04070c5806 --- /dev/null +++ b/crates/mir-importer/src/translator/cute/smem_mma_emit.rs @@ -0,0 +1,1183 @@ +/* + * 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::{Located, 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::emit::{load_through, recover_cta_shared_pointer, struct_field_addr}; +use 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, +}; + +pub(super) 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); + } +} + +pub(super) 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)) +} + +#[allow(clippy::too_many_arguments)] +pub(super) 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) +} + +pub(super) 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:?}" + ))) + }) +} + +pub(super) 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) + ))) + }) +} + +pub(super) 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:?}" + ))) + }) +} + +pub(super) 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)) +} + +pub(super) 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)) +} + +pub(super) 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)) +} + +fn plan(facts: &TiledMmaRust) -> CuteTiledMmaPlanAttr { + CuteTiledMmaPlanAttr::mxf4_128x128x128(facts.placement.clone()) +} + +fn fixed_plan() -> CuteTiledMmaPlanAttr { + CuteTiledMmaPlanAttr::mxf4_128x128x128(canonical_data_placement()) +} + +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, + ) +} + +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" + ))) + }) +} + +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 Mxfp4TiledMma" + ))) + }) +} + +#[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)) +} + +#[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) +} + +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) +} + +pub(super) 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}" + ))) + }) +} + +#[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", + ) +} + +#[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 Mxfp4TiledMma".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", + ) +} + +#[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", + ) +} + +#[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", + ) +} + +#[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", + ) +} + +#[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", + ) +} + +#[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", + ) +} + +#[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/crates/mir-importer/src/translator/cute/static_config.rs b/crates/mir-importer/src/translator/cute/static_config.rs new file mode 100644 index 0000000000..d8efd348e0 --- /dev/null +++ b/crates/mir-importer/src/translator/cute/static_config.rs @@ -0,0 +1,508 @@ +/* + * 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 +//! L, C<8>>, T2, C<1>>> +//! │ └────── shape ──────┘ └───── stride ─────┘ +//! └─ exact DefId: cute_rs::markers::L +//! ``` +//! +//! 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, +}; + +const C: &str = "cute_rs::markers::C"; +const T2: &str = "cute_rs::markers::T2"; +const LAYOUT: &str = "cute_rs::markers::L"; +const ROW_MAJOR: &str = "cute_rs::markers::RowMajor"; +const COL_MAJOR: &str = "cute_rs::markers::ColMajor"; +const CP_ASYNC: &str = "cute_rs::markers::CpAsync"; +const LEADING_DIM: &str = "cute_rs::markers::LeadingDim"; +const SWIZZLE: &str = "cute_rs::markers::Swizzle"; +const COMPOSED: &str = "cute_rs::markers::Composed"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct CopyAtom { + pub bytes: u32, +} + +#[derive(Debug, Clone)] +pub(crate) struct CooperativeConfig { + pub atom: CopyAtom, + pub thread_layout: Layout, + pub value_layout: Layout, + pub tile_layout: Layout, + 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. + #[allow(dead_code)] + pub element_type: Ty, + // Keep the validated byte width beside the type so emission cannot re-guess it. + #[allow(dead_code)] + pub element_bytes: i64, +} + +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) +} + +fn canonical_adt_path(def: &AdtDef) -> String { + let crate_name = def.krate().name.to_string(); + let mut segments = Vec::new(); + let mut current = Some(def.def_id()); + while let Some(def_id) = current { + let printed = def_id.name(); + let segment = printed.as_str().rsplit("::").next().unwrap_or_default(); + segments.push(segment.to_owned()); + current = def_id.parent(); + } + super::canonical_path_from_leaf_segments(&crate_name, segments) +} + +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) +} + +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(", ") + )) + } +} + +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")) +} + +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")) +} + +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")), + } +} + +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")), + } +} + +fn decode_tuple(ty: &Ty) -> Result { + let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { + return Err(format!("layout tuple marker must be an ADT, 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", + )?)) + } + T2 => { + expect_schema(&args, &["Type", "Type"], T2)?; + Ok(IntTuple::Tuple(vec![ + decode_tuple(&type_arg(&args, 0, T2)?)?, + decode_tuple(&type_arg(&args, 1, T2)?)?, + ])) + } + found => Err(format!("expected `{C}` or `{T2}`, found `{found}`")), + } +} + +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()) + }), + } +} + +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.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) +} + +fn decode_atom(ty: &Ty) -> Result { + let args = adt_owned(ty, CP_ASYNC)?; + expect_schema(&args, &["Const"], CP_ASYNC)?; + let bytes = u32::try_from(unsigned_const( + &const_arg(&args, 0, CP_ASYNC)?, + "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 }) +} + +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) +} + +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") + }) +} + +pub(crate) fn decode_smem_layout(ty: &Ty) -> Result { + if let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() + && canonical_adt_path(&def) == COMPOSED + { + expect_schema(&args, &["Type", "Const", "Type"], COMPOSED)?; + let outer = decode_swizzle(&type_arg(&args, 0, COMPOSED)?)?; + let offset = signed_const(&const_arg(&args, 1, COMPOSED)?, "composed offset")?; + let inner = 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(outer, offset, inner, 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::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.inner().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 { + pub smem_layout: ComposedLayout, + 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, +} + +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::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::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, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[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/crates/mir-importer/src/translator/cute/tensor.rs b/crates/mir-importer/src/translator/cute/tensor.rs new file mode 100644 index 0000000000..eb4f884eba --- /dev/null +++ b/crates/mir-importer/src/translator/cute/tensor.rs @@ -0,0 +1,243 @@ +/* + * 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. + +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, +}; + +const TENSOR: &str = "cute_rs::tensor::Tensor"; +const TENSOR_MUT: &str = "cute_rs::tensor::TensorMut"; +const CONTIGUOUS_1D: &str = "cute_rs::tensor::Contiguous1D"; +const ZIPPED_1D: &str = "cute_rs::tensor::Zipped1D"; +const TILE_1D: &str = "cute_rs::tensor::Tile1D"; +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 { + pub logical: Ty, + pub storage: Ty, + pub access: CuteTensorAccessAttr, + pub layout: CuteTensorLayoutAttr, +} + +/// Rust-side facts for `RegisterTile`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct RegisterTileRust { + pub element: Ty, + pub size: u64, +} + +pub(super) fn canonical_adt_path(def: &AdtDef) -> String { + let crate_name = def.krate().name.to_string(); + let mut segments = Vec::new(); + let mut current = Some(def.def_id()); + while let Some(def_id) = current { + let printed = def_id.name(); + let segment = printed.as_str().rsplit("::").next().unwrap_or_default(); + segments.push(segment.to_owned()); + current = def_id.parent(); + } + super::canonical_path_from_leaf_segments(&crate_name, segments) +} + +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(", ") + )) + } +} + +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")), + } +} + +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")), + } +} + +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")) +} + +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/crates/mir-importer/src/translator/cute/tensor_emit.rs b/crates/mir-importer/src/translator/cute/tensor_emit.rs new file mode 100644 index 0000000000..db35fe6e65 --- /dev/null +++ b/crates/mir-importer/src/translator/cute/tensor_emit.rs @@ -0,0 +1,710 @@ +/* + * 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::{Located, 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::tensor::{decode_register_tile, decode_tensor_view}; + +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); + } +} + +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)) +} + +#[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) +} + +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) +} + +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)) +} + +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))) +} + +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)) +} + +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)) +} + +#[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", + ) +} + +#[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 destination_ty = destination.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read zipped_divide destination type: {error:?}" + ))) + })?; + let result = decode_tensor_view(&destination_ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid zipped_divide destination: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "zipped_divide destination is not a v0 tensor view".to_string() + )) + })?; + let tile_size = result.layout.tile_size().ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "zipped_divide result does not carry a tile width".to_string() + )) + })?; + let (tensor, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + 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, after.or(prev)); + let result = op.deref(ctx).get_result(0); + finish_value( + ctx, + op, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "tensor_zipped_divide", + ) +} + +#[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 slice = CuteTensorSliceOp::new(ctx, tensor, tile_index); + let op = slice.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, + "tensor_slice", + ) +} + +#[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" + }, + ) +} + +#[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") +} + +#[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") +} + +#[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/crates/mir-importer/src/translator/mod.rs b/crates/mir-importer/src/translator/mod.rs index 784ca58758..8d11a85ed0 100644 --- a/crates/mir-importer/src/translator/mod.rs +++ b/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/crates/mir-importer/src/translator/rvalue/fn_ptr.rs b/crates/mir-importer/src/translator/rvalue/fn_ptr.rs index 7a6e0de1dc..0e80108d47 100644 --- a/crates/mir-importer/src/translator/rvalue/fn_ptr.rs +++ b/crates/mir-importer/src/translator/rvalue/fn_ptr.rs @@ -28,6 +28,10 @@ use std::num::NonZeroUsize; /// name, never 0 so it cannot look like a null pointer) and casts it /// int -> ptr. See the comment at the `Rvalue::Cast` arm for why a token /// stands in for a code address on the device. +fn is_cute_rs_crate(crate_name: &str) -> bool { + matches!(crate_name, "cute_rs" | "cute-rs") +} + pub(super) fn translate_reify_fn_pointer( ctx: &mut Context, body: &mir::Body, @@ -56,6 +60,16 @@ pub(super) fn translate_reify_fn_pointer( )) ); }; + let crate_name = fn_def.krate().name; + if is_cute_rs_crate(crate_name.as_str()) { + let item_name = fn_def.name(); + return input_err!( + loc, + TranslationErr::unsupported(format!( + "CuTe function `{item_name}` must be called directly; converting it to a function pointer would erase its layout configuration" + )) + ); + } let raw_intrinsic = crate::translator::terminator::intrinsics::generated::require_supported_raw_intrinsic( fn_def, &loc, diff --git a/crates/mir-importer/src/translator/rvalue/place_read.rs b/crates/mir-importer/src/translator/rvalue/place_read.rs index d1e6a0f392..6d18929f91 100644 --- a/crates/mir-importer/src/translator/rvalue/place_read.rs +++ b/crates/mir-importer/src/translator/rvalue/place_read.rs @@ -80,6 +80,11 @@ pub fn translate_place( prev_op: Option>, loc: Location, ) -> TranslationResult<(Value, Option>)> { + if place.projection.is_empty() + && let Some(value) = value_map.get_direct_value(place.local) + { + return Ok((value, prev_op)); + } match classify_place_read_strategy(ctx, place, value_map)? { PlaceReadStrategy::Address => { if let Some((value, last_op)) = translate_place_load_from_address( diff --git a/crates/mir-importer/src/translator/terminator/mod.rs b/crates/mir-importer/src/translator/terminator/mod.rs index 34f8d9889c..33270597e9 100644 --- a/crates/mir-importer/src/translator/terminator/mod.rs +++ b/crates/mir-importer/src/translator/terminator/mod.rs @@ -968,6 +968,25 @@ fn translate_call( let (pattern_name, call_name, substs_args, type_substs) = extract_func_info(func, &loc)?; let loc = call_debug_location(pattern_name.as_deref(), loc); + // Keep CuTe calls as CuTe operations while their tensor and layout + // meaning is still visible. Ordinary source helpers continue through the + // normal Rust-body path. + if let Some(result) = crate::translator::cute::try_translate_cute_call( + ctx, + body, + func, + args, + destination, + &target.map(|block| block), + block_ptr, + prev_op, + value_map, + block_map, + &loc, + ) { + return result; + } + // Is the trait-method Self type SharedArray? Shared with // `values::classify_call` so intrinsic dispatch and destination-slot // classification can't drift. @@ -2088,7 +2107,10 @@ fn extract_function_item_target( requires_direct_dispatch: fn_def.is_intrinsic() || instance.is_foreign_item() || !instance.has_body() - || matches!(crate_name.as_str(), "cuda_device" | "cuda-device" | "libm") + || matches!( + crate_name.as_str(), + "cuda_device" | "cuda-device" | "libm" | "cute_rs" | "cute-rs" + ) || generated_direct_call_only, })) } diff --git a/crates/mir-importer/src/translator/types.rs b/crates/mir-importer/src/translator/types.rs index ee5c23b8e7..d2ac860e83 100644 --- a/crates/mir-importer/src/translator/types.rs +++ b/crates/mir-importer/src/translator/types.rs @@ -662,6 +662,29 @@ pub fn translate_type( translate_pointer_like(ctx, &pointee, origin) } rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Adt(adt_def, substs)) => { + if let Some(result) = + crate::translator::cute::pipeline::try_translate_tma_load_pipeline_type( + ctx, rust_ty, + ) + { + return result; + } + if let Some(result) = + crate::translator::cute::scheduler::try_translate_work_tile_type(ctx, rust_ty) + { + return result; + } + if let Some(result) = + crate::translator::cute::block_scaled::try_translate_block_scaled_type(ctx, rust_ty) + { + return result; + } + if let Some(result) = + crate::translator::cute::tensor::try_translate_tensor_view_type(ctx, rust_ty) + { + return result; + } + // Get the trimmed name (just the type name without path) let trimmed_name = adt_def.trimmed_name(); diff --git a/crates/mir-importer/src/translator/values.rs b/crates/mir-importer/src/translator/values.rs index 450665aef0..35f975f1a0 100644 --- a/crates/mir-importer/src/translator/values.rs +++ b/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/crates/rustc-codegen-cuda/Cargo.lock b/crates/rustc-codegen-cuda/Cargo.lock index b7075b19c9..260533a397 100644 --- a/crates/rustc-codegen-cuda/Cargo.lock +++ b/crates/rustc-codegen-cuda/Cargo.lock @@ -206,6 +206,7 @@ name = "cuda-oxide-codegen" version = "0.1.0" dependencies = [ "cuda-target-spec", + "dialect-cute", "dialect-iket", "dialect-mir", "dialect-nvvm", @@ -226,6 +227,24 @@ dependencies = [ 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" @@ -498,6 +517,8 @@ version = "0.2.1" dependencies = [ "combine", "cuda-oxide-codegen", + "cute-layout", + "dialect-cute", "dialect-iket", "dialect-mir", "dialect-nvvm", diff --git a/crates/rustc-codegen-cuda/src/collector.rs b/crates/rustc-codegen-cuda/src/collector.rs index c361689734..2b7de4cbef 100644 --- a/crates/rustc-codegen-cuda/src/collector.rs +++ b/crates/rustc-codegen-cuda/src/collector.rs @@ -235,6 +235,191 @@ enum CollectDecision { Forbidden { crate_name: String, fn_path: String }, } +/// CuTe entry points whose bodies are compiler-owned stubs. +/// +/// Ordinary `cute-rs` helpers are still collected. Only these exact defining +/// paths are replaced with `cute.*` operations at their direct call sites. +const CUTE_RECOGNITION_STUB_PATHS: &[&str] = &[ + "::tensor::make_tensor_read", + "::tensor::make_tensor_write", + "::tensor::zipped_divide_read", + "::tensor::zipped_divide_write", + "::tensor::slice_read", + "::tensor::slice_write", + "::tensor::tensor_tile_is_full", + "::tensor::tensor_tile_base", + "::tensor::tensor_load_tile", + "::tensor::tensor_store_tile", + "::tensor::tensor_store_element_abs", + "::block_scaled::__compiler::block_scaled_make", + "::block_scaled::__compiler::block_scaled_thread_row", + "::block_scaled::__compiler::block_scaled_k_tile", + "::block_scaled::__compiler::block_scaled_load_k64", + "::block_scaled::__compiler::block_scaled_dot_k64", + "::scheduler::__compiler::scheduler_new_1d", + "::scheduler::__compiler::scheduler_has_work", + "::scheduler::__compiler::scheduler_current_tile", + "::scheduler::__compiler::work_tile_coordinates", + "::scheduler::__compiler::scheduler_advance", + "::pipeline::__compiler::tma_load_pipeline_from_raw_base", + "::pipeline::__compiler::tma_load_pipeline_init", + "::pipeline::__compiler::pipeline_state_new_producer", + "::pipeline::__compiler::pipeline_state_new_consumer", + "::pipeline::__compiler::pipeline_state_slot", + "::pipeline::__compiler::pipeline_state_advance", + "::pipeline::__compiler::pipeline_producer_acquire", + "::pipeline::__compiler::pipeline_producer_expect_tx", + "::pipeline::__compiler::pipeline_consumer_wait", + "::pipeline::__compiler::pipeline_consumer_release", + "::pipeline::__compiler::pipeline_producer_tail", + "::tiled_copy::__compiler::shared_tensor_overlay", + "::block_scaled_mma::__compiler::tiled_mma_slice", + "::block_scaled_mma::__compiler::fragment_fill", + "::block_scaled_mma::__compiler::mma_load_scales", + "::block_scaled_mma::__compiler::fragment_slice_k", + "::block_scaled_mma::__compiler::mma_load_a", + "::block_scaled_mma::__compiler::mma_partition_b", + "::block_scaled_mma::__compiler::tiled_gemm", + "::epilogue::__compiler::epilogue_smem_overlay", + "::epilogue::__compiler::epilogue_warp_slice", + "::block_scaled_mma::__compiler::epilogue_store_fragment", + "::epilogue::__compiler::epilogue_sync_reusable", + "::epilogue::__compiler::epilogue_sync_ready_for_tma", + "::epilogue::__compiler::epilogue_half", + "::pipeline::__compiler::tma_store_producer_acquire", + "::pipeline::__compiler::tma_store_producer_commit", + "::pipeline::__compiler::tma_store_producer_tail", + "::tile::load_tile", + "::tile::store_tile", + "::tile::assume_div", + "::cooperative::copy_g2s", + "::mma::load_matrix_a", + "::mma::load_matrix_b", + "::tma::copy_tma_2d", + "::tma::copy_tma_s2g_2d", +]; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum CuteImplOwner { + Tensor, + TensorMut, + BlockScaledTensor, + BlockScaledThreadRow, + BlockScaledTile64, + LoadedBlockScaledTile64, + StaticPersistentTileScheduler, + WorkTile, + TmaLoadPipeline, + PipelineState, + SharedTensor, + Mxfp4TiledMma, + Mxf4AccumulatorTile2x8, + Mxf4ScaleTile128, + Sm120Epilogue128x128, + Sm120EpilogueWarp128x128, + TmaStorePipeline, +} + +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; + } + if CUTE_RECOGNITION_STUB_PATHS.contains(&definition_path) { + return true; + } + + let Some(impl_owner) = impl_owner else { + return false; + }; + let module = match impl_owner { + CuteImplOwner::Tensor | CuteImplOwner::TensorMut => "tensor", + CuteImplOwner::BlockScaledTensor + | CuteImplOwner::BlockScaledThreadRow + | CuteImplOwner::BlockScaledTile64 + | CuteImplOwner::LoadedBlockScaledTile64 => "block_scaled", + CuteImplOwner::StaticPersistentTileScheduler | CuteImplOwner::WorkTile => "scheduler", + CuteImplOwner::TmaLoadPipeline + | CuteImplOwner::PipelineState + | CuteImplOwner::TmaStorePipeline => "pipeline", + CuteImplOwner::SharedTensor => "tiled_copy", + CuteImplOwner::Mxfp4TiledMma + | CuteImplOwner::Mxf4AccumulatorTile2x8 + | CuteImplOwner::Mxf4ScaleTile128 => "block_scaled_mma", + CuteImplOwner::Sm120Epilogue128x128 => "epilogue", + CuteImplOwner::Sm120EpilogueWarp128x128 => "block_scaled_mma", + }; + 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" | "slice" | "is_full" | "base" | "load" + ), + CuteImplOwner::TensorMut => matches!( + method, + "from_disjoint_slice" | "zipped_divide" | "slice" | "store" | "store_linear" + ), + 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" | "current_tile" | "advance") + } + CuteImplOwner::WorkTile => method == "coordinates", + CuteImplOwner::TmaLoadPipeline => 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" | "slot" | "advance"), + CuteImplOwner::SharedTensor => method == "from_raw_parts", + CuteImplOwner::Mxfp4TiledMma => matches!( + method, + "get_slice" | "load_scale_atom_128" | "load_a_128" | "get_b_tile_k64" + ), + CuteImplOwner::Mxf4AccumulatorTile2x8 => matches!(method, "zero" | "accumulate_k64"), + 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::TmaStorePipeline => { + 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 @@ -1747,6 +1932,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) { @@ -1912,6 +2104,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 +2424,52 @@ 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, _) => match self.tcx.item_name(def.did()).to_string().as_str() { + "Tensor" => Some(CuteImplOwner::Tensor), + "TensorMut" => Some(CuteImplOwner::TensorMut), + "BlockScaledTensor" => Some(CuteImplOwner::BlockScaledTensor), + "BlockScaledThreadRow" => Some(CuteImplOwner::BlockScaledThreadRow), + "BlockScaledTile64" => Some(CuteImplOwner::BlockScaledTile64), + "LoadedBlockScaledTile64" => Some(CuteImplOwner::LoadedBlockScaledTile64), + "StaticPersistentTileScheduler" => { + Some(CuteImplOwner::StaticPersistentTileScheduler) + } + "WorkTile" => Some(CuteImplOwner::WorkTile), + "TmaLoadPipeline" => Some(CuteImplOwner::TmaLoadPipeline), + "PipelineState" => Some(CuteImplOwner::PipelineState), + "SharedTensor" => Some(CuteImplOwner::SharedTensor), + "Mxfp4TiledMma" => Some(CuteImplOwner::Mxfp4TiledMma), + "Mxf4AccumulatorTile2x8" => Some(CuteImplOwner::Mxf4AccumulatorTile2x8), + "Mxf4ScaleTile128" => Some(CuteImplOwner::Mxf4ScaleTile128), + "Sm120Epilogue128x128" => Some(CuteImplOwner::Sm120Epilogue128x128), + "Sm120EpilogueWarp128x128" => Some(CuteImplOwner::Sm120EpilogueWarp128x128), + "TmaStorePipeline" => Some(CuteImplOwner::TmaStorePipeline), + _ => None, + }, + _ => None, + } + }); + is_cute_recognition_stub_identity(crate_name.as_str(), &definition_path, impl_owner) + } + /// 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( @@ -2748,9 +2999,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 +3122,180 @@ 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_linear", + Some(CuteImplOwner::TensorMut), + )); + 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::{impl#17}::from_slices", + Some(CuteImplOwner::BlockScaledTensor), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::block_scaled::{impl#19}::k_tile", + Some(CuteImplOwner::BlockScaledThreadRow), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::block_scaled::{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}::coordinates", + Some(CuteImplOwner::WorkTile), + )); + 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::TmaLoadPipeline), + )); + 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::TmaLoadPipeline), + ), + "the rejected split-ring constructor must not be body-translated" + ); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::tiled_copy::__compiler::shared_tensor_overlay", + None, + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::tiled_copy::{impl#3}::from_raw_parts", + Some(CuteImplOwner::SharedTensor), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::block_scaled_mma::__compiler::tiled_gemm", + None, + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::block_scaled_mma::{impl#7}::load_a_128", + Some(CuteImplOwner::Mxfp4TiledMma), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::block_scaled_mma::{impl#4}::zero", + Some(CuteImplOwner::Mxf4AccumulatorTile2x8), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::block_scaled_mma::{impl#5}::pairs_at_unchecked", + Some(CuteImplOwner::Mxf4ScaleTile128), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::epilogue::__compiler::epilogue_sync_ready_for_tma", + None, + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::epilogue::{impl#2}::tma_half", + Some(CuteImplOwner::Sm120Epilogue128x128), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::block_scaled_mma::{impl#9}::store_tile", + Some(CuteImplOwner::Sm120EpilogueWarp128x128), + )); + 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::TmaStorePipeline), + )); + + // 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::{impl#21}::value_pair", + Some(CuteImplOwner::LoadedBlockScaledTile64), + )); + assert!(!is_cute_recognition_stub_identity( + "cute_rs", + "::scheduler::{impl#4}::current", + 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", + "::block_scaled_mma::{impl#7}::load_b_pair", + Some(CuteImplOwner::Mxfp4TiledMma), + )); + assert!(!is_cute_recognition_stub_identity( + "cute_rs", + "::epilogue::{impl#3}::store_atom", + Some(CuteImplOwner::Sm120EpilogueWarp128x128), + )); + assert!(!is_cute_recognition_stub_identity( + "other", + "::tensor::make_tensor_read", + None, + )); + } + #[test] fn kernel_def_paths_are_entry_points() { // Bare and fully-qualified kernel names: the marker is the final segment. diff --git a/cuteir/README.md b/cuteir/README.md new file mode 100644 index 0000000000..fa17f4d7e4 --- /dev/null +++ b/cuteir/README.md @@ -0,0 +1,89 @@ +# CuTe-style GPU kernels in Rust + +`cuteir` lets a Rust kernel describe GPU data by its meaning, not only by its +address. + +The two basic ideas are small: + +- a **Tensor** is storage plus a layout; +- a **Layout** says where a logical coordinate lives in that storage. + +That is enough to express operations such as “split this vector into +four-value tiles” or “copy this matrix tile into shared memory.” The compiler +can then choose how to lower those operations without asking the kernel author +to rewrite the pointer arithmetic by hand. + +```text +plain storage + layout + │ + ▼ + Tensor ──► tile ──► slice ──► copy / add / matrix multiply +``` + +## Semantic frontend + +An **IR** is the compiler's working form of a program. A **dialect** is a named +set of operations in that IR. A **backend** turns the IR into code another +compiler or the GPU can consume. + +The frontend keeps CuTe operations in `dialect-cute` while they are still +recognizable: + +```text +Rust kernel using cute-rs +Tensor → zipped_divide → slice → copy / add / matrix multiply + │ + ▼ + high-level dialect-cute + │ + ▼ + preparation + verification + │ + ▼ + selected backend continuation +``` + +The whole-module verifier checks that tensor provenance, scheduler and +pipeline state, TMA transactions, shared-memory MMA, and epilogue operations +form a complete semantic story. It does not choose or execute a backend. + +## What lives here + +| Path | What it does | +| :--- | :--- | +| [`cute-layout`](cute-layout) | Layout math shared by the Rust API and compiler. | +| [`cute-rs`](cute-rs) | Device-facing Tensor, Layout, copy, pipeline, and matrix multiply-and-accumulate (MMA) types. | +| [`dialect-cute`](dialect-cute) | CuTe operations and backend-neutral whole-module verification in Pliron. | +| [`elementwise_cute`](examples/elementwise_cute) | Adds two vectors through per-thread Tensor tiles. | +| [`nvfp4_gemv_cute`](examples/nvfp4_gemv_cute) | Multiplies a packed FP4 matrix by a packed FP4 vector. | +| [`blockscale_gemm_cute`](examples/blockscale_gemm_cute) | Multiplies two block-scaled FP4 matrices with the Tensor Memory Accelerator (TMA) and tensor cores. | + +The examples are deliberately ordered from smallest to largest: + +```text +elementwise GEMV GEMM +vector tiles → scaled row tiles → global/shared matrix tiles +load + add load + dot TMA + pipeline + MMA + epilogue +``` + +## Run the examples + +Run these commands from the repository root after completing the normal +CUDA Oxide setup: + +```bash +cargo oxide run elementwise_cute + +cargo oxide run nvfp4_gemv_cute --arch sm_120a -- \ + --m 512 --k 256 --l 1 + +cargo oxide run blockscale_gemm_cute --arch sm_120a +``` + +The elementwise example uses ordinary `f32` and `f16` operations. The two FP4 +examples use Blackwell instructions and should be built for `sm_120a`; their +own READMEs give the exact device and shape requirements. + +Each program checks every GPU output bit pattern against a deterministic host +result before reporting success. Backend-specific artifact and code-shape +checks live with the backend implementation rather than in this shared layer. diff --git a/cuteir/cute-layout/Cargo.toml b/cuteir/cute-layout/Cargo.toml new file mode 100644 index 0000000000..6fd31e22c6 --- /dev/null +++ b/cuteir/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/cuteir/cute-layout/src/algebra.rs b/cuteir/cute-layout/src/algebra.rs new file mode 100644 index 0000000000..b557b3ba61 --- /dev/null +++ b/cuteir/cute-layout/src/algebra.rs @@ -0,0 +1,1248 @@ +/* + * 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? +//! 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. + +use alloc::format; +use alloc::string::String; +use alloc::vec; +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(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 {} + +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 + .leaves() + .into_iter() + .zip(l.stride.leaves()) + .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()), + }, + } +} + +/// 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))) +} + +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`] 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`] 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`] 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)) +} + +/// 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. +pub fn zipped_divide_by_shape(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`] 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) +} + +/// 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)) +} + +/// 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.leaves(); + let strides = layout.stride.leaves(); + 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 + .leaves() + .into_iter() + .fold(0, gcd_nonnegative); + Ok(gcd_nonnegative(vector_size, rest_stride_gcd).max(1)) +} + +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 +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::string::ToString; + + 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() + } + + #[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"); + } + + /// 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" + ); + } + + #[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}" + ); + } + } + } + + /// 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_by_shape(&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})" + ); + } + } + } + } + } + + /// 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); + } + + #[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::>()); + } + + #[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}"); + } + } + + #[test] + fn max_alignment_is_conservative_for_negative_strides() { + assert_eq!(max_alignment(&l("8:-1")).unwrap(), 1); + } +} diff --git a/cuteir/cute-layout/src/composed.rs b/cuteir/cute-layout/src/composed.rs new file mode 100644 index 0000000000..b0e95c0c78 --- /dev/null +++ b/cuteir/cute-layout/src/composed.rs @@ -0,0 +1,534 @@ +/* + * 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 +//! │ +//! ▼ +//! inner layout ──> plain offset ──> + static offset ──> outer swizzle +//! │ +//! ▼ +//! physical tile offset +//! +//! call(c) = outer(offset + inner(c)) +//! ``` +//! +//! The names “inner” and “outer” describe evaluation order: the inner layout +//! runs first; the outer swizzle runs last. 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 +//! inner layout: (2,4):(4,1) +//! row mode: extent 2, stride 4 +//! column mode: extent 4, stride 1 +//! fixed offset: 0 +//! outer 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 inner 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(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 `outer(offset + inner(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 { + outer: Swizzle, + offset: i64, + inner: Layout, + unit: OffsetUnit, +} + +impl ComposedLayout { + /// Build the mapping `outer(offset + inner(coordinate))`. + /// + /// `offset`, every stride in `inner`, and the input expected by `outer` + /// 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( + outer: Swizzle, + offset: i64, + inner: Layout, + unit: OffsetUnit, + ) -> Result { + if !outer.is_valid() { + return Err(ComposedLayoutError::InvalidSwizzle); + } + Ok(ComposedLayout { + outer, + offset, + inner, + 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 + /// inner layout (2,3):(3,1), coordinate (1,2) + /// inner offset 5 -> identity swizzle -> final offset 5 + /// ``` + pub fn from_layout(inner: Layout, unit: OffsetUnit) -> Self { + ComposedLayout { + outer: Swizzle::IDENTITY, + offset: 0, + inner, + unit, + } + } + + /// Return the swizzle applied last. + /// + /// “Outer” describes composition order; it does not mean an outer tile or + /// an absolute address. + pub const fn outer(&self) -> Swizzle { + self.outer + } + + /// Return the fixed displacement added after the inner layout and before + /// the swizzle, measured in [`Self::unit`]. + pub const fn offset(&self) -> i64 { + self.offset + } + + /// Return the plain coordinate-to-offset layout evaluated first. + pub const fn inner(&self) -> &Layout { + &self.inner + } + + /// 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 inner layout. + /// + /// Swizzling only permutes locations, so it does not change this count. + pub fn size(&self) -> i64 { + self.inner.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 inner layout; + /// compiler input should use [`Self::checked_call`] instead. The module's + /// [concrete example](crate::composed#concrete-example) shows the complete + /// inner-layout, fixed-offset, and swizzle calculation. + pub fn call(&self, coord: &IntTuple) -> i64 { + self.call_with_inner_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_inner_delta(coord, 0) + } + + /// Map a position shortly after the item selected by `coord`. + /// + /// `delta` is added to the inner 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 `CpAsync<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_inner_delta(&self, coord: &IntTuple, delta: i64) -> i64 { + self.outer + .apply(self.offset + self.inner.call(coord) + delta) + } + + /// Checked form of [`Self::call_with_inner_delta`]. + /// + /// Returns `None` for an invalid coordinate or arithmetic overflow. + pub fn checked_call_with_inner_delta(&self, coord: &IntTuple, delta: i64) -> Option { + let input = self + .offset + .checked_add(self.inner.checked_call(coord)?)? + .checked_add(delta)?; + Some(self.outer.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 inner 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 + .outer + .base + .checked_add(added_bits) + .ok_or(ComposedLayoutError::ArithmeticOverflow)?; + let outer = Swizzle::try_new(self.outer.bits, new_base, self.outer.shift) + .ok_or(ComposedLayoutError::InvalidSwizzle)?; + let offset = self + .offset + .checked_mul(element_bytes) + .ok_or(ComposedLayoutError::ArithmeticOverflow)?; + let stride = checked_scale_tuple(&self.inner.stride, element_bytes) + .ok_or(ComposedLayoutError::ArithmeticOverflow)?; + + Ok(ComposedLayout { + outer, + offset, + inner: Layout::new(self.inner.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 + /// inner layout's alignment + /// fixed offset + /// swizzle's chunk-alignment factor + /// │ + /// ▼ + /// gcd(inner alignment, fixed offset, swizzle alignment) + /// ``` + /// + /// Example: an inner 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 inner_alignment = max_alignment(&self.inner)?; + // gcd(x, 0) = x, so a zero offset drops out on its own. + let with_offset = gcd(inner_alignment, unsigned_to_i64(self.offset.unsigned_abs())); + if self.outer.bits == 0 { + return Ok(with_offset.max(1)); + } + Ok(gcd(with_offset, self.outer.max_alignment()).max(1)) + } +} + +impl fmt::Display for ComposedLayout { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "{} o {} o {} [{}]", + self.outer, self.offset, self.inner, self.unit + ) + } +} + +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), + } +} + +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) +} + +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) +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::string::ToString; + + fn l(text: &str) -> Layout { + text.parse().unwrap() + } + + #[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.outer.apply(8 + layout.inner.call(&coord)) + ); + } + assert_eq!(layout.to_string(), "S<2,0,3> o 8 o 4:1 [bytes]"); + } + + #[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.outer(), Swizzle::new(2, 4, 3)); + assert_eq!(bytes.offset(), 32); + assert_eq!(bytes.inner().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); + } + } + + #[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) + ); + } + + #[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 { + bits: 4, + base: 3, + shift: 3, + }; + assert_eq!( + ComposedLayout::new(invalid, 0, l("4:1"), OffsetUnit::Bytes), + Err(ComposedLayoutError::InvalidSwizzle) + ); + } + + #[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); + } + + #[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 inner_limited = ComposedLayout::new( + Swizzle::new(3, 4, 4), + 0, + l("(8,16):(1,4)"), + OffsetUnit::Bytes, + ) + .unwrap(); + assert_eq!(inner_limited.max_alignment().unwrap(), 4); + } + + #[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/cuteir/cute-layout/src/lib.rs b/cuteir/cute-layout/src/lib.rs new file mode 100644 index 0000000000..31611ab498 --- /dev/null +++ b/cuteir/cute-layout/src/lib.rs @@ -0,0 +1,877 @@ +/* + * 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. +//! +//! # 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. `CpAsync<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 +//! CpAsync<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. + +#![no_std] + +extern crate alloc; + +pub mod algebra; +pub mod composed; +pub mod predication; +pub mod swizzle; +pub mod validation; + +pub use composed::{ComposedLayout, ComposedLayoutError, OffsetUnit}; +pub use predication::{ + AtomCoord, GmemMatrixBounds, PredicatedSource, PredicationError, PredicationResult, + SafeSourceOffset, plan_uniform_cp_async_source, +}; +pub use swizzle::Swizzle; +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, +}; + +use alloc::vec; +use alloc::vec::Vec; +use core::fmt; +use core::str::FromStr; + +/// 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(i64), + /// An ordered group of integers or further groups. + Tuple(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`] 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. + /// + /// 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`] requires congruent shape and stride tuples so every shape + /// number has a matching stride number in the same position. + pub fn 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.congruent(y)) + } + _ => false, + } + } + + /// Returns every integer from left to right and discards parentheses. + /// + /// ```text + /// input: ((2,3),4) + /// output: [2, 3, 4] parentheses are gone + /// ``` + /// + /// Algebra operations use this when they need the numbers but not the + /// nested [mode](crate#modes) structure. + pub fn leaves(&self) -> Vec { + let mut out = Vec::new(); + self.collect_leaves(&mut out); + out + } + + 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. + 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()) + } + } + } +} + +/// 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, + } +} + +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.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.leaves(); + 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. + 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()], + } + } + + /// 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()), + }, + } + } +} + +// 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 {} + +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.congruent(&stride) { + return Err(ParseLayoutError); + } + Ok(Layout { shape, stride }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::vec; + + fn t(v: &[i64]) -> IntTuple { + IntTuple::Tuple(v.iter().map(|x| IntTuple::Leaf(*x)).collect()) + } + + #[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); + } + + #[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); + } + + #[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 + } + + #[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); + } + + #[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/cuteir/cute-layout/src/predication.rs b/cuteir/cute-layout/src/predication.rs new file mode 100644 index 0000000000..a9d9b8c2df --- /dev/null +++ b/cuteir/cute-layout/src/predication.rs @@ -0,0 +1,479 @@ +/* + * 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 module computes data only: the compile-time half of the split. The +//! predicate's value depends on runtime numbers, so the emitted kernel +//! evaluates it; what is settled here is which comparisons exist and why +//! every branch-free outcome is memory-safe: +//! +//! ```text +//! this module (compile time) emitted kernel (GPU runtime) +//! ────────────────────────── ──────────────────────────── +//! which comparisons must exist evaluate them: row < rows? +//! prove base + 0 is a safe fallback pick the offset with select +//! reject pitch < cols, overflow, issue cp.async with the chosen +//! bad atom widths, loudly source_bytes (16 / 8 / 0) +//! ``` +//! +//! Later, the compiler must use the plan to choose the offset and source +//! size consistently for every participating thread, and then issue +//! `cp.async` for each one. + +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(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 }, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + const MATRIX: GmemMatrixBounds = GmemMatrixBounds { + rows: 3, + columns: 10, + leading_dim: 12, + base_alignment_bytes: 16, + }; + + #[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 }, + }) + ); + } + + #[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 }, + }) + ); + } + + #[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 }, + }) + ); + } + } + + #[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, + }) + ); + } + + #[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, + }) + ); + } + + #[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/cuteir/cute-layout/src/swizzle.rs b/cuteir/cute-layout/src/swizzle.rs new file mode 100644 index 0000000000..690ba13065 --- /dev/null +++ b/cuteir/cute-layout/src/swizzle.rs @@ -0,0 +1,386 @@ +/* + * 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 inner 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 { bits: B, base: M, shift: S }` represents CuTe's `S`. +/// 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 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 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 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(bits: u32, base: u32, shift: i32) -> Self { + assert!( + Self::parameters_are_valid(bits, base, shift), + "invalid CuTe swizzle" + ); + Swizzle { bits, base, 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(bits: u32, base: u32, shift: i32) -> Option { + if Self::parameters_are_valid(bits, base, shift) { + Some(Swizzle { bits, base, shift }) + } else { + None + } + } + + /// Return whether `S` has disjoint fields inside the supported low + /// 32-bit swizzle window. + /// + /// `bits`, `base`, and `shift` mean `B`, `M`, and `S`, respectively. + pub const fn parameters_are_valid(bits: u32, base: u32, shift: i32) -> bool { + let distance = shift.unsigned_abs(); + if distance < bits { + return false; + } + let Some(base_and_bits) = base.checked_add(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.bits, self.base, self.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.bits) - 1; + mask << (self.base as i64 + self.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.bits) - 1; + mask << (self.base as i64 + (-self.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.shift >= 0 { + masked >> self.shift + } else { + masked << (-self.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.shift >= 0 { + y >> self.shift + } else { + y << (-self.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.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.shift < 0 { + return None; + } + self.base.checked_add(self.shift as u32) + } +} + +impl fmt::Display for Swizzle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "S<{},{},{}>", self.bits, self.base, self.shift) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[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); + } + + #[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()); + } + + #[test] + #[should_panic(expected = "invalid CuTe swizzle")] + fn unchecked_construction_fails_loudly() { + let _ = Swizzle::new(4, 3, 3); + } + + #[test] + fn identity_is_identity() { + for x in 0..256 { + assert_eq!(Swizzle::IDENTITY.apply(x), x); + } + } + + #[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::>()); + } + + #[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" + ); + } + } + + #[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); + } + + #[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/cuteir/cute-layout/src/validation.rs b/cuteir/cute-layout/src/validation.rs new file mode 100644 index 0000000000..88cf84cf4d --- /dev/null +++ b/cuteir/cute-layout/src/validation.rs @@ -0,0 +1,1668 @@ +/* + * 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. When the compiler emits the real copy +//! instructions, it must still check facts known only while the kernel runs: +//! valid pointers, matrix bounds, pitch (the element distance between row +//! starts), edge handling, the actual shared allocation, and the current +//! thread index. + +use alloc::collections::BTreeMap; +use alloc::format; +use alloc::string::String; +use alloc::vec; +use alloc::vec::Vec; +use core::fmt; + +use crate::algebra::make_layout_tv; +use crate::{ComposedLayout, IntTuple, Layout, OffsetUnit}; + +/// 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; + +/// 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. + /// + /// Before issuing the copy, the compiler must prove that the runtime + /// capacity is at least this value. It 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(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 {} + +/// 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(&'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 {} + +fn plan_error(message: impl Into) -> CooperativePlanError { + CooperativePlanError(message.into()) +} + +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.leaves() { + 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(()) +} + +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_inner_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. Before using it, the compiler must still check facts known only +/// while the kernel runs: enough shared storage, matrix bounds and row spacing, +/// valid and aligned pointers, the current thread index, and how many bytes an +/// edge copy should read before zero-filling the rest. +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.inner(), "shared-memory inner")?; + 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 inner", smem_layout.inner()), + ] { + 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, + }) +} + +/// 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`] 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 +/// ``` +/// +/// `CpAsync` 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 `CpAsync`, 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 + .inner() + .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_inner_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(()) +} + +fn tv_coord(thread: i64, value: i64) -> IntTuple { + IntTuple::Tuple(vec![IntTuple::Leaf(thread), IntTuple::Leaf(value)]) +} + +fn positive_count(name: &'static str, value: i64) -> ValidationResult { + if value <= 0 { + Err(ValidationError::InvalidCount { name, value }) + } else { + Ok(value) + } +} + +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 }) +} + +/// 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(()) +} + +/// 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.inner().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.leaves(); + let col_leaves = modes[1].stride.leaves(); + 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.outer(); + if swizzle.bits == 0 { + return Ok(()); + } + let m_bytes = i64::from(swizzle.base) + element_bytes.trailing_zeros() as i64; + if swizzle.bits > 3 || m_bytes != 4 || swizzle.shift != 3 { + return Err(ValidationError::InvalidCount { + name: "swizzle is not one of TMA's encodable byte-unit S modes", + value: i64::from(swizzle.bits), + }); + } + let pitch_bytes = columns * element_bytes; + if pitch_bytes != (16 << swizzle.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.outer(); + if swizzle.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.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", + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Swizzle; + use alloc::string::ToString; + + fn l(text: &str) -> Layout { + text.parse().unwrap() + } + + 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) + } + + #[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); + } + + #[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}"); + } + + #[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}"); + } + + #[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 { .. }) + )); + } + + #[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 { .. }) + )); + } + + #[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(); + } + + #[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", .. }) + )); + } + + #[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(); + } + + #[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(); + } + + #[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 { .. }) + )); + } + + #[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 { .. }) + )); + } + + #[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" + )) + ); + } + + #[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)); + } + + #[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)); + } + + #[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)); + } + + #[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/cuteir/cute-rs/Cargo.toml b/cuteir/cute-rs/Cargo.toml new file mode 100644 index 0000000000..eec6a8a1d8 --- /dev/null +++ b/cuteir/cute-rs/Cargo.toml @@ -0,0 +1,20 @@ +[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] +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] +trybuild = "1" diff --git a/cuteir/cute-rs/src/block_scaled.rs b/cuteir/cute-rs/src/block_scaled.rs new file mode 100644 index 0000000000..b3da2a61cb --- /dev/null +++ b/cuteir/cute-rs/src/block_scaled.rs @@ -0,0 +1,860 @@ +/* + * 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. + +use core::marker::PhantomData; + +use crate::cooperative::GmemMatrix; +use crate::markers::LeadingDim; +use crate::numeric::{E2M1, Mxf4E2M1, PackedE2M1x2, UE8M0, UE8M0x2, UE8M0x4}; +use crate::tiled_copy::{GlobalCopyTensor, SharedTensor}; +use crate::{Tensor, TensorElement, assume_div, load_tile}; + +/// 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; + +/// Rows whose K values are adjacent in memory. +/// +/// ```text +/// row 0: [k0 k1 k2 ...] +/// row 1: [k0 k1 k2 ...] +/// ``` +/// +/// `Mode` is either [`Mkl`] or [`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 { + rows: usize, + k: usize, + 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 { + rows: usize, + groups: usize, + 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, + } + } + + /// Return the byte offset for `(row, K group)`. + #[must_use] + #[inline(always)] + pub const fn offset(&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 + } +} + +/// The two [`UE8M0`] 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`]) separate from B scales ([`Nkl`]) so the MMA +/// call cannot swap them accidentally. +#[repr(transparent)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct MmaScalePair { + bits: u32, + 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 { + bits: u32, + 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, + } + } +} + +/// 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 Mxfp4BlockScaledTensor<'a, Mode: Copy> { + values: Tensor<'a, Mxf4E2M1, KMajor, f16>, + scales: Tensor<'a, UE8M0, DenseBlockScaleKMajor, u8>, +} + +impl<'a, Mode: Copy> Mxfp4BlockScaledTensor<'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.offset(row, pair * 2); + MmaScalePair { + bits: (self.scales.storage[offset] as u32) + | ((self.scales.storage[offset + 1] as u32) << 8), + mode: PhantomData, + } + } +} + +/// 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 { + rest_m: usize, + rest_k: usize, +} + +// Map one logical scale coordinate into Blackwell's 512-byte 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) + }}; +} + +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), + } + } + + /// Return the byte offset for `(batch, row, K group)`. + #[must_use] + #[inline(always)] + pub const fn offset(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 + } +} + +/// 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, +> { + values: Tensor<'a, Data, DataLayout, u8>, + scales: Tensor<'a, Scale, Sm1xxBlockScaleKMajor, u8>, +} + +impl<'a, Mode: Copy> BlockScaledTensor<'a, E2M1, UE8M0, 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) + } +} + +/// 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> { + values: &'a [u8], + scales: &'a [u8], + value_row_base: usize, + scale_row_base: usize, + 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> { + values: &'a [u8], + scales: &'a [u8], + value_base: usize, + scale_base: usize, + 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 { + values_lo: [u8; 16], + values_hi: [u8; 16], + scale0: f32, + scale1: f32, + scale2: f32, + scale3: f32, + mode: PhantomData, +} + +impl LoadedBlockScaledTile64 { + /// Return packed FP4 pair `PAIR` in increasing K order. + #[must_use] + #[inline(always)] + pub fn value_pair(self) -> PackedE2M1x2 { + const { assert!(PAIR < 32) }; + let bits = if PAIR < 16 { + self.values_lo[PAIR] + } else { + self.values_hi[PAIR - 16] + }; + PackedE2M1x2::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) + } +} + +/// Stable calls recognized by the CuTe importer. +/// +/// Kernel authors use the methods above. Keeping these functions in one +/// private module gives the compiler an exact boundary that does not depend +/// on whether rustc inlines an ergonomic method wrapper. +mod __compiler { + use super::*; + + #[inline(never)] + pub(super) const fn block_scaled_make<'a, Mode: Copy>( + values: &'a [u8], + scales: &'a [u8], + rows: usize, + k: usize, + ) -> BlockScaledTensor<'a, E2M1, UE8M0, 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)] + pub(super) fn block_scaled_thread_row<'a, Mode: Copy>( + tensor: BlockScaledTensor<'a, E2M1, UE8M0, 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)] + 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)] + 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_div::<16>(tile.value_base) }; + let scale_base = unsafe { assume_div::<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)] + pub(super) fn block_scaled_dot_k64( + lhs: LoadedBlockScaledTile64, + rhs: LoadedBlockScaledTile64, + mut acc: f32, + ) -> f32 { + 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 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn canonical_scale_layout_matches_blackwell_golden_offsets() { + let layout = Sm1xxBlockScaleKMajor::<16>::new(256, 128); + + assert_eq!(layout.offset(0, 0, 0), 0); + assert_eq!(layout.offset(0, 31, 0), 496); + assert_eq!(layout.offset(0, 32, 0), 4); + assert_eq!(layout.offset(0, 127, 3), 511); + assert_eq!(layout.offset(0, 0, 4), 512); + assert_eq!(layout.offset(0, 128, 0), 1_024); + assert_eq!(layout.offset(1, 0, 0), 2_048); + assert_eq!(layout.storage_len(2), 4_096); + } + + #[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); + } + + #[test] + fn mxf4_tensor_keeps_logical_k_and_materializes_packed_carriers() { + let values = [f16::from_bits(0); 128]; + let scales = [127u8; 16]; + let matrix = Mxfp4BlockScaledTensor::::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); + } + + #[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); + + #[repr(align(4))] + struct AlignedScales([u8; 32]); + let values = [f16::from_bits(0); 256]; + let scales = AlignedScales([127; 32]); + let matrix = Mxfp4BlockScaledTensor::::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); + } + + #[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.offset(32, 0), 128); + assert_eq!(dense.storage_len(), 512); + assert_eq!(canonical.offset(0, 32, 0), 4); + assert_eq!(canonical.storage_len(1), 512); + } + + #[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.offset(0, row, 0)); + for group in 0..CANONICAL_SCALE_K_GROUPS { + assert_eq!(word * 4 + group, canonical.offset(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/cuteir/cute-rs/src/block_scaled_mma.rs b/cuteir/cute-rs/src/block_scaled_mma.rs new file mode 100644 index 0000000000..e124affb0f --- /dev/null +++ b/cuteir/cute-rs/src/block_scaled_mma.rs @@ -0,0 +1,965 @@ +/* + * 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. + +use core::marker::PhantomData; + +use crate::block_scaled::{Mkl, MmaScalePair, Nkl, ScalePack4, SharedScaleAtom, Sm120ScaleAtom}; +use crate::epilogue::Sm120EpilogueWarp128x128; +use crate::markers::ColMajor; +use crate::mma::{AccC, FragA, FragB, load_matrix_a}; +use crate::numeric::{Mxf4E2M1, UE8M0x4}; +use crate::tiled_copy::SharedTensor; + +/// 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 { + a: ScalePack4, + b0: ScalePack4, + b1: ScalePack4, + b2: ScalePack4, + 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 { + a0: ScalePack4, + a1: ScalePack4, + b0: ScalePack4, + b1: ScalePack4, + b2: ScalePack4, + b3: ScalePack4, + b4: ScalePack4, + b5: ScalePack4, + b6: ScalePack4, + 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 { + c00: AccC, + c01: AccC, + c02: AccC, + c03: AccC, + c04: AccC, + c05: AccC, + c06: AccC, + c07: AccC, + c10: AccC, + c11: AccC, + c12: AccC, + c13: AccC, + c14: AccC, + c15: AccC, + c16: AccC, + c17: AccC, +} + +/// 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::accumulate_k64`], 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> { + base: *mut f16, + capacity: usize, + warp_n: usize, + k_half: usize, + 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)] + unsafe fn load(&self, mma: &Mxfp4TiledMma) -> (FragB, FragB) { + 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. + /// + /// # 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 accumulate_k64( + &mut self, + mma: &Mxfp4TiledMma, + a: (FragA, FragA), + b: Mxf4BTileK64<'_, SharedLayout>, + scales: Mxf4ScalePairs128, + ) { + unsafe { __compiler::tiled_gemm(self, mma, a, b, scales) }; + } +} + +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) }; + } +} + +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) } + } +} + +// Map `(warp M position, first-or-second M block)` to a 16-row block number. +#[inline(always)] +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)] +const fn b_pair_atom_128(warp_n: usize, pair_n: usize) -> usize { + warp_n + 2 * pair_n +} + +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) }, + } + } +} + +/// 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 Mxfp4TiledMma { + lane: u32, + layout: PhantomData, +} + +impl Mxfp4TiledMma { + /// 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::accumulate_k64`]. `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), + ) -> FragA { + 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), + ) -> (FragB, FragB) { + let loaded = + unsafe { load_matrix_a::(source.carrier(), warp_tile, self.lane) }; + ( + FragB([loaded.0[0], loaded.0[2]]), + FragB([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, + ) -> (FragA, FragA) { + 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, + ) -> (FragB, FragB) { + 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`] + /// 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`] 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: FragA, + scale_a: MmaScalePair, + b: FragB, + scale_b: MmaScalePair, + c: AccC, + ) -> AccC { + let _ = (self, a, scale_a, b, scale_b, c); + unreachable!("cute-rs block-scaled MMA executed outside recognized CuTe device compilation") + } +} + +/// Stable compiler boundaries for the SM120 shared-memory MMA path. +/// +/// Public methods above keep the kernel readable. These exact calls let the +/// importer retain the same tensor, fragment, and tiled-MMA meaning whether +/// rustc inlines those wrappers or leaves them as calls. +#[doc(hidden)] +pub mod __compiler { + use super::*; + + #[inline(never)] + pub const fn tiled_mma_slice(lane: u32) -> Mxfp4TiledMma { + Mxfp4TiledMma { + lane, + layout: PhantomData, + } + } + + #[inline(never)] + pub const fn fragment_fill(fill: f32) -> Mxf4AccumulatorTile2x8 { + let cell = AccC([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, + } + } + + #[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); + } + } + + #[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) }, + } + } + + #[inline(never)] + pub unsafe fn mma_load_scales( + mma: &Mxfp4TiledMma, + 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))), + } + } + } + + #[inline(never)] + pub unsafe fn mma_load_a( + mma: &Mxfp4TiledMma, + source: &SharedTensor, + warp_m: usize, + k_half: usize, + ) -> (FragA, FragA) { + 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)), + ) + } + } + + #[inline(never)] + pub const fn mma_partition_b<'a, SharedLayout>( + _mma: &Mxfp4TiledMma, + 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, + } + } + + #[inline(never)] + pub unsafe fn tiled_gemm( + accumulators: &mut Mxf4AccumulatorTile2x8, + mma: &Mxfp4TiledMma, + a: (FragA, FragA), + 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) }; + } + } +} + +#[cfg(test)] +mod tests { + use super::{Mxf4AccumulatorTile2x8, Mxf4BTileK64, Mxfp4TiledMma, a_atom_128, b_pair_atom_128}; + use crate::block_scaled::{Mkl, Nkl, SharedScaleAtom, Sm120ScaleAtom}; + use crate::mma::AccC; + + #[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::() + ); + } + + #[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]] + ); + } + + #[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 = Mxfp4TiledMma::<()>::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 = Mxfp4TiledMma::<()>::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/cuteir/cute-rs/src/cooperative.rs b/cuteir/cute-rs/src/cooperative.rs new file mode 100644 index 0000000000..91072a1547 --- /dev/null +++ b/cuteir/cute-rs/src/cooperative.rs @@ -0,0 +1,148 @@ +/* + * 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; + +/// 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 full +/// pointer, size, pitch, and layout contract. +#[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, +} + +/// 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") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[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/cuteir/cute-rs/src/epilogue.rs b/cuteir/cute-rs/src/epilogue.rs new file mode 100644 index 0000000000..4df6fa00c7 --- /dev/null +++ b/cuteir/cute-rs/src/epilogue.rs @@ -0,0 +1,375 @@ +/* + * 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 +//! [`AccC`]. 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::markers::{Composed, RowMajor, Swizzle}; +use crate::{AccC, SmemTile}; + +/// 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 = Composed, 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::(); + +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: *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::*; + + #[inline(never)] + pub const unsafe fn epilogue_smem_overlay(base: *mut f16) -> Sm120Epilogue128x128 { + Sm120Epilogue128x128 { base } + } + + #[inline(never)] + pub const unsafe fn epilogue_warp_slice( + tile: Sm120Epilogue128x128, + warp_id: usize, + lane: u32, + ) -> Sm120EpilogueWarp128x128 { + Sm120EpilogueWarp128x128 { + base: tile.base, + warp_id, + lane, + } + } + + #[inline(never)] + pub unsafe fn epilogue_sync_reusable(tile: Sm120Epilogue128x128) { + let _ = tile; + unreachable!("cute-rs epilogue sync executed outside recognized device compilation") + } + + #[inline(never)] + pub unsafe fn epilogue_sync_ready_for_tma(tile: Sm120Epilogue128x128) { + let _ = tile; + unreachable!("cute-rs epilogue publication executed outside recognized device compilation") + } + + #[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: *mut f16, + warp_id: usize, + 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: AccC) { + 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") + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::markers::ReifySmem2D; + use cuda_device::mma_frag; + + #[test] + fn half_layout_is_exactly_one_b128_tma_tile() { + assert_eq!(::ROWS, 128); + assert_eq!(::COLS, 64); + assert_eq!(::ROW_STRIDE, 64); + assert_eq!(::COL_STRIDE, 1); + assert_eq!(::SWIZZLE_B, 3); + assert_eq!(::SWIZZLE_M, 3); + assert_eq!(::SWIZZLE_S, 3); + assert_eq!(SM120_EPILOGUE_HALF_BYTES, 16 * 1024); + assert_eq!(Sm120Epilogue128x128::BYTES, 32 * 1024); + } + + #[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)); + } + + #[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); + } + + #[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/cuteir/cute-rs/src/lib.rs b/cuteir/cute-rs/src/lib.rs new file mode 100644 index 0000000000..067fe522e4 --- /dev/null +++ b/cuteir/cute-rs/src/lib.rs @@ -0,0 +1,99 @@ +/* + * 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. +//! +//! The types describe what a kernel means. The compiler turns a small set of +//! recognized functions into GPU operations: +//! +//! ```text +//! Tensor + Layout + TiledCopy +//! │ +//! ▼ +//! Rust kernel code +//! │ +//! ▼ +//! load / copy / MMA / store +//! │ +//! ▼ +//! selected GPU backend +//! ``` +//! +//! 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_div`] 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 block_scaled; +pub mod block_scaled_mma; +pub mod cooperative; +pub mod epilogue; +pub mod markers; +pub mod mma; +pub mod numeric; +pub mod pipeline; +pub mod scheduler; +pub mod tensor; +pub mod tile; +pub mod tiled_copy; +pub mod tma; + +pub use block_scaled::{ + BlockScaledTensor, BlockScaledThreadRow, BlockScaledTile64, DenseBlockScaleKMajor, KMajor, + LoadedBlockScaledTile64, Mkl, MmaScalePair, Mxfp4BlockScaledTensor, Nkl, ScalePack4, + SharedScaleAtom, Sm1xxBlockScaleKMajor, Sm120ScaleAtom, +}; +pub use block_scaled_mma::{ + Mxf4AccumulatorTile2x8, Mxf4BTileK64, Mxf4ScalePairs, Mxf4ScalePairs128, Mxf4ScaleStage, + Mxf4ScaleTile128, Mxfp4TiledMma, +}; +pub use cooperative::{GmemMatrix, SmemTile, copy_g2s}; +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, Composed, CpAsync, L, LeadingDim, RowMajor, Swizzle, T2}; +pub use mma::{AccC, FragA, FragB, load_matrix_a, load_matrix_b}; +pub use numeric::{E2M1, Mxf4E2M1, PackedE2M1x2, UE8M0, UE8M0x2, UE8M0x4}; +pub use pipeline::{Consumer, PipelineState, Producer, TmaLoadPipeline, TmaStorePipeline}; +pub use scheduler::{StaticPersistentTileScheduler, WorkTile}; +pub use tensor::{Contiguous1D, RegisterTile, Tensor, TensorElement, TensorMut, Tile1D, Zipped1D}; +pub use tile::{assume_div, load_tile, store_tile}; +pub use tiled_copy::{GlobalCopyTensor, SharedTensor, TiledCopy}; +pub use tma::{TmaDesc, copy_tma_2d, copy_tma_s2g_2d, tile_bytes}; diff --git a/cuteir/cute-rs/src/markers.rs b/cuteir/cute-rs/src/markers.rs new file mode 100644 index 0000000000..5db9f90c2f --- /dev/null +++ b/cuteir/cute-rs/src/markers.rs @@ -0,0 +1,217 @@ +/* + * 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: +//! +//! ```text +//! 32 rows × 8 columns, row-major +//! +//! L< T2, C<8>>, T2, C<1>> > +//! └── shape ─────┘ └─ strides ───┘ +//! row: 8 +//! column: 1 +//! ``` + +use core::marker::PhantomData; + +/// One integer stored in a type, usually a size or stride. +pub struct C; + +/// Two type-level values grouped together. +/// +/// Pairs can be nested. For example, `(A, B), C` and `A, (B, C)` describe +/// different layouts. +pub struct T2(PhantomData<(A, B)>); + +/// A layout made from a `Shape` and its `Stride`. +/// +/// A stride is the number of elements to skip when a coordinate increases by +/// one. +pub struct L(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 `cp.async` copy of `BYTES` bytes from global to shared memory. +pub struct CpAsync; + +/// 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 `Inner`, adding `OFFSET`, then applying +/// `Outer`. +/// +/// ```text +/// coordinate ── Inner ──► element offset +/// + OFFSET +/// │ +/// ▼ +/// Outer ──► final address +/// ``` +/// +/// `OFFSET` and `Inner` count elements here. The compiler converts the whole +/// rule to bytes after it knows the element size. +pub struct Composed(PhantomData<(Outer, Inner)>); + +/// 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 +/// ``` +pub trait ReifySmem2D { + /// Number of rows. + const ROWS: i64; + /// Number of columns. + const COLS: i64; + /// Number of elements skipped when the row increases by one. + const ROW_STRIDE: i64; + /// Number of elements skipped when the column increases by one. + const COL_STRIDE: i64; + /// Fixed starting offset in elements; `0` for a plain layout. + const OFFSET: i64; + /// Swizzle settings in element units; `B = 0` means no swizzle. + const SWIZZLE_B: u32; + /// Number of low address bits left unchanged by the swizzle. + const SWIZZLE_M: u32; + /// Distance between the two address fields used by the swizzle. + const SWIZZLE_S: i32; +} + +impl ReifySmem2D + for L, C>, T2, C>> +{ + const ROWS: i64 = R; + const COLS: i64 = CC; + const ROW_STRIDE: i64 = SR; + const COL_STRIDE: i64 = SC; + const OFFSET: i64 = 0; + const SWIZZLE_B: u32 = 0; + const SWIZZLE_M: u32 = 4; + const SWIZZLE_S: i32 = 3; +} + +impl ReifySmem2D for RowMajor { + const ROWS: i64 = R; + const COLS: i64 = CC; + const ROW_STRIDE: i64 = CC; + const COL_STRIDE: i64 = 1; + const OFFSET: i64 = 0; + const SWIZZLE_B: u32 = 0; + const SWIZZLE_M: u32 = 4; + const SWIZZLE_S: i32 = 3; +} + +impl ReifySmem2D for ColMajor { + const ROWS: i64 = R; + const COLS: i64 = CC; + const ROW_STRIDE: i64 = 1; + const COL_STRIDE: i64 = R; + const OFFSET: i64 = 0; + const SWIZZLE_B: u32 = 0; + const SWIZZLE_M: u32 = 4; + const SWIZZLE_S: i32 = 3; +} + +/// Read one composed layout. Do not nest `Composed`: this implementation +/// returns the outer swizzle and would hide the inner one. +impl ReifySmem2D + for Composed, OFFSET, Inner> +{ + const ROWS: i64 = Inner::ROWS; + const COLS: i64 = Inner::COLS; + const ROW_STRIDE: i64 = Inner::ROW_STRIDE; + const COL_STRIDE: i64 = Inner::COL_STRIDE; + const OFFSET: i64 = OFFSET; + const SWIZZLE_B: u32 = B; + const SWIZZLE_M: u32 = M; + const SWIZZLE_S: i32 = S; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn markers_are_compile_time_only() { + type Shape = T2, C<8>>; + type Stride = T2, C<1>>; + type Affine = L; + type Smem = Composed, 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/cuteir/cute-rs/src/mma.rs b/cuteir/cute-rs/src/mma.rs new file mode 100644 index 0000000000..79cdf6eeb1 --- /dev/null +++ b/cuteir/cute-rs/src/mma.rs @@ -0,0 +1,104 @@ +/* + * 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: FragA [u32; 4], FragB [u32; 2], AccC [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. +//! +use crate::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 FragA(pub [u32; 4]); + +/// One thread's part of a 16x8 `f16` B tile: two 32-bit registers. +#[repr(transparent)] +#[derive(Clone, Copy)] +pub struct FragB(pub [u32; 2]); + +/// 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 AccC(pub [f32; 4]); + +impl AccC { + /// Four zero values for starting `C = A*B + C`. + pub const ZERO: AccC = AccC([0.0; 4]); +} + +/// Load one 16x16 `f16` A tile into the 32 threads of a warp. +/// +/// Each thread receives its [`FragA`]. `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, +) -> FragA { + 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 [`FragB`]. `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, +) -> FragB { + let _ = (src, warp_tile, lane); + unreachable!("cute-rs `load_matrix_b` executed outside device compilation") +} diff --git a/cuteir/cute-rs/src/numeric.rs b/cuteir/cute-rs/src/numeric.rs new file mode 100644 index 0000000000..3092cadd8d --- /dev/null +++ b/cuteir/cute-rs/src/numeric.rs @@ -0,0 +1,240 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Small number formats used by block-scaled kernels. +//! +//! 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 +//! E2M1 u8 holds 2 FP4 values +//! UE8M0 u8 holds 1 scale +//! Mxf4E2M1 f16 holds 4 FP4 values +//! UE8M0x4 u32 holds 4 scales +//! ``` +//! +//! [`PackedE2M1x2`] 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 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 E2M1; + +/// 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 UE8M0; + +/// Four neighboring [`UE8M0`] 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 [`E2M1`] 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 E2M1 {} +impl TensorElement for UE8M0 {} +impl TensorElement for Mxf4E2M1 {} +impl TensorElement for UE8M0x4 {} + +/// Two [`E2M1`] 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 PackedE2M1x2(u8); + +impl PackedE2M1x2 { + /// 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 [`UE8M0`] 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(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), + ) + } +} + +#[inline(always)] +fn e2m1_to_f32(nibble: u8) -> f32 { + const MAGNITUDES: [f32; 8] = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]; + let magnitude = MAGNITUDES[(nibble & 0x07) as usize]; + if nibble & 0x08 == 0 { + magnitude + } else { + -magnitude + } +} + +#[inline(always)] +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), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn e2m1_pair_preserves_all_bits() { + for bits in 0u8..=u8::MAX { + assert_eq!(PackedE2M1x2::from_bits(bits).to_bits(), bits); + } + } + + #[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); + } + + #[test] + fn e2m1_conversion_covers_every_code_bit_exactly() { + 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() + ); + } + } + + #[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()); + } +} diff --git a/cuteir/cute-rs/src/pipeline.rs b/cuteir/cute-rs/src/pipeline.rs new file mode 100644 index 0000000000..fb41a6cb49 --- /dev/null +++ b/cuteir/cute-rs/src/pipeline.rs @@ -0,0 +1,555 @@ +/* + * 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. This matches CUTLASS +//! `PipelineTmaAsync`. + +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 { + slot: u32, + phase: u32, + 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 shared-buffer slot. + #[must_use] + #[inline(always)] + pub fn slot(&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. +/// +/// - `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 TmaLoadPipeline { + full: *mut Barrier, + 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. This matches CUTLASS `PipelineTmaStore`. +#[derive(Clone, Copy, Debug, Default)] +pub struct TmaStorePipeline; + +impl TmaStorePipeline { + /// 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 + TmaLoadPipeline +{ + /// 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::TmaLoadPipeline; + /// const _: usize = TmaLoadPipeline::<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::*; + + #[inline(never)] + pub fn tma_store_producer_acquire() { + let _ = TmaStorePipeline::::MAX_PENDING; + unreachable!("cute-rs TMA-store acquire executed outside recognized device compilation") + } + + #[inline(never)] + pub fn tma_store_producer_commit() { + let _ = TmaStorePipeline::::MAX_PENDING; + unreachable!("cute-rs TMA-store commit executed outside recognized device compilation") + } + + #[inline(never)] + pub fn tma_store_producer_tail() { + let _ = TmaStorePipeline::::MAX_PENDING; + unreachable!("cute-rs TMA-store tail executed outside recognized device compilation") + } + + #[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, + } + } + + #[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, + } + } + + #[inline(never)] + pub fn pipeline_state_slot( + state: &PipelineState, + ) -> usize { + state.slot as usize + } + + #[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; + } + } + + #[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, + ) -> TmaLoadPipeline { + let _ = TmaLoadPipeline::::STORAGE_BYTES; + TmaLoadPipeline { + full: base, + empty: unsafe { base.add(STAGES) }, + } + } + + #[inline(never)] + pub unsafe fn tma_load_pipeline_init< + const STAGES: usize, + const CONSUMER_WARPS: u32, + const TX_BYTES: u32, + >( + pipeline: TmaLoadPipeline, + init_thread: u32, + ) { + let _ = TmaLoadPipeline::::STORAGE_BYTES; + let _ = (pipeline.full, pipeline.empty, init_thread); + unreachable!("cute-rs TMA-load init executed outside recognized device compilation") + } + + #[inline(never)] + pub unsafe fn pipeline_producer_acquire< + const STAGES: usize, + const CONSUMER_WARPS: u32, + const TX_BYTES: u32, + >( + pipeline: TmaLoadPipeline, + state: &PipelineState, + ) { + let _ = TmaLoadPipeline::::STORAGE_BYTES; + let _ = (pipeline, state); + unreachable!("cute-rs producer acquire executed outside recognized device compilation") + } + + #[inline(never)] + pub unsafe fn pipeline_producer_expect_tx< + const STAGES: usize, + const CONSUMER_WARPS: u32, + const TX_BYTES: u32, + >( + pipeline: TmaLoadPipeline, + state: &PipelineState, + ) -> *mut Barrier { + let _ = TmaLoadPipeline::::STORAGE_BYTES; + let _ = (pipeline, state); + unreachable!("cute-rs producer expect-tx executed outside recognized device compilation") + } + + #[inline(never)] + pub unsafe fn pipeline_consumer_wait< + const STAGES: usize, + const CONSUMER_WARPS: u32, + const TX_BYTES: u32, + >( + pipeline: TmaLoadPipeline, + state: &PipelineState, + ) { + let _ = TmaLoadPipeline::::STORAGE_BYTES; + let _ = (pipeline, state); + unreachable!("cute-rs consumer wait executed outside recognized device compilation") + } + + #[inline(never)] + pub unsafe fn pipeline_consumer_release< + const STAGES: usize, + const CONSUMER_WARPS: u32, + const TX_BYTES: u32, + >( + pipeline: TmaLoadPipeline, + state: &PipelineState, + ) { + let _ = TmaLoadPipeline::::STORAGE_BYTES; + let _ = (pipeline, state); + unreachable!("cute-rs consumer release executed outside recognized device compilation") + } + + #[inline(never)] + pub unsafe fn pipeline_producer_tail< + const STAGES: usize, + const CONSUMER_WARPS: u32, + const TX_BYTES: u32, + >( + pipeline: TmaLoadPipeline, + state: PipelineState, + ) { + let _ = TmaLoadPipeline::::STORAGE_BYTES; + let _ = (pipeline, state); + unreachable!("cute-rs producer tail executed outside recognized device compilation") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn producer_and_consumer_use_opposite_initial_phases() { + let producer = PipelineState::::new(); + let consumer = PipelineState::::new(); + assert_eq!((producer.slot(), producer.phase()), (0, 1)); + assert_eq!((consumer.slot(), consumer.phase()), (0, 0)); + } + + #[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 (slot, phase) in expected { + assert_eq!((state.slot(), state.phase()), (slot, phase)); + state.advance(); + } + } + + #[test] + fn state_is_exactly_two_device_register_values() { + assert_eq!(size_of::>(), 8); + assert_eq!(size_of::>(), 8); + } + + #[test] + fn static_contract_sizes_two_aligned_barrier_rings() { + type Pipeline = TmaLoadPipeline<3, 8, 32_768>; + const BYTES: usize = Pipeline::STORAGE_BYTES; + assert_eq!(Pipeline::STORAGE_ALIGN, 8); + assert_eq!(BYTES, 6 * size_of::()); + } + + #[test] + fn tma_store_pipeline_is_zero_storage_and_uses_stage_minus_one() { + assert_eq!(size_of::>(), 0); + assert_eq!(TmaStorePipeline::<1>::MAX_PENDING, 0); + assert_eq!(TmaStorePipeline::<4>::MAX_PENDING, 3); + } +} diff --git a/cuteir/cute-rs/src/scheduler.rs b/cuteir/cute-rs/src/scheduler.rs new file mode 100644 index 0000000000..9c68b468e5 --- /dev/null +++ b/cuteir/cute-rs/src/scheduler.rs @@ -0,0 +1,413 @@ +/* + * 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 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 WorkTile { + linear: usize, + m_tile: usize, + n_tile: usize, + batch: usize, +} + +impl + WorkTile +{ + /// Return `(linear, M tile, N tile, batch)` together. + /// + /// Kernels normally unpack this once at the start of a tile iteration. + #[must_use] + #[inline(always)] + pub fn coordinates(self) -> (usize, usize, usize, usize) { + __compiler::work_tile_coordinates(self) + } + + /// Tile number when M changes first, then N, then batch. + #[must_use] + #[inline(always)] + pub const fn linear(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`] 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, +> { + current_linear: usize, + 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, + }) + } + } + + /// Return the current tile coordinates, or `None` when this block is done. + #[must_use] + #[inline(always)] + pub fn current(&self) -> Option> { + let linear = self.current_linear; + if linear >= Self::TOTAL_TILES { + return None; + } + + // M changes first, then N, then batch. Compute each remainder with + // multiply/subtract so LLVM can reuse the matching division result. + 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; + + Some(WorkTile { + linear, + m_tile, + n_tile, + batch, + }) + } + + /// True while this block still has an output tile to process. + #[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 split keeps the bounds check outside the coordinate arithmetic in + /// the kernel loop. + #[must_use] + #[inline(always)] + pub fn current_tile(&self) -> WorkTile { + __compiler::scheduler_current_tile(self) + } + + /// Move to this block's next output tile. + #[inline(always)] + pub fn advance(&mut self) { + __compiler::scheduler_advance(self); + } + + /// Choose the number of blocks for a one-dimensional persistent launch. + /// + /// 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 resident_grid_size(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, WorkTile}; + + #[inline(never)] + pub fn scheduler_new_1d() + -> StaticPersistentTileScheduler { + let _ = StaticPersistentTileScheduler::::TOTAL_TILES; + unreachable!("cute-rs scheduler creation executed outside recognized device compilation") + } + + #[inline(never)] + pub fn scheduler_has_work( + scheduler: &StaticPersistentTileScheduler, + ) -> bool { + scheduler.current_linear + < StaticPersistentTileScheduler::::TOTAL_TILES + } + + #[inline(never)] + pub fn scheduler_current_tile< + const M_TILES: usize, + const N_TILES: usize, + const BATCHES: usize, + >( + scheduler: &StaticPersistentTileScheduler, + ) -> WorkTile { + 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; + WorkTile { + linear, + m_tile, + n_tile, + batch, + } + } + + #[inline(never)] + pub fn work_tile_coordinates< + const M_TILES: usize, + const N_TILES: usize, + const BATCHES: usize, + >( + tile: WorkTile, + ) -> (usize, usize, usize, usize) { + (tile.linear, tile.m_tile, tile.n_tile, tile.batch) + } + + #[inline(never)] + pub fn scheduler_advance( + scheduler: &mut StaticPersistentTileScheduler, + ) { + scheduler.current_linear = scheduler.current_linear.saturating_add(scheduler.stride); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use core::mem::size_of; + + #[test] + fn m_fastest_mapping_crosses_n_and_batch_boundaries() { + 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(); + let tile = scheduler.current().unwrap(); + assert_eq!( + (tile.linear(), tile.m_tile(), tile.n_tile(), tile.batch()), + (linear, m_tile, n_tile, batch) + ); + } + } + + #[test] + fn grid_stride_170_covers_the_full_16_by_16_problem_once() { + type Scheduler = StaticPersistentTileScheduler<16, 16, 1>; + 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 let Some(tile) = scheduler.current() { + let linear = tile.linear(); + 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(); + } + } + + assert!(seen.into_iter().all(|visited| visited)); + } + + #[test] + fn uneven_final_wave_has_the_expected_worker_lengths() { + 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.current().is_some() { + count += 1; + scheduler.advance(); + } + assert_eq!(count, if worker < 86 { 2 } else { 1 }); + } + } + + #[test] + fn worker_constructor_rejects_invalid_physical_grids() { + 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()); + } + + #[test] + fn resident_grid_is_checked_and_capped() { + type Scheduler = StaticPersistentTileScheduler<16, 16, 1>; + assert_eq!(Scheduler::resident_grid_size(170, 1), Some(170)); + assert_eq!(Scheduler::resident_grid_size(170, 2), Some(256)); + assert_eq!(Scheduler::resident_grid_size(512, 8), Some(256)); + assert_eq!(Scheduler::resident_grid_size(0, 1), None); + assert_eq!(Scheduler::resident_grid_size(170, 0), None); + + type TooWide = StaticPersistentTileScheduler<{ u32::MAX as usize }, 2, 1>; + assert_eq!(TooWide::resident_grid_size(u32::MAX, u32::MAX), None); + } + + #[test] + fn completed_scheduler_stays_invalid() { + type Scheduler = StaticPersistentTileScheduler<1, 1, 1>; + let mut scheduler = Scheduler::for_worker(0, 1).unwrap(); + assert!(scheduler.current().is_some()); + scheduler.advance(); + assert!(scheduler.current().is_none()); + scheduler.advance(); + assert!(scheduler.current().is_none()); + } + + #[test] + fn scheduler_state_is_only_current_and_stride() { + type Scheduler = StaticPersistentTileScheduler<16, 16, 1>; + assert_eq!(size_of::(), 2 * size_of::()); + assert_eq!(Scheduler::TOTAL_TILES, 256); + } +} diff --git a/cuteir/cute-rs/src/tensor.rs b/cuteir/cute-rs/src/tensor.rs new file mode 100644 index 0000000000..ae1ce711c7 --- /dev/null +++ b/cuteir/cute-rs/src/tensor.rs @@ -0,0 +1,673 @@ +/* + * 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. + +use core::marker::PhantomData; +use core::ops::Add; + +use cuda_device::DisjointSlice; + +use crate::{assume_div, load_tile, store_tile}; + +/// Says which Rust storage type may hold one or more logical elements. +/// +/// For example, `E2M1: 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 {} + +/// 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 { + 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 { + 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 { + base: usize, + 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([T; N]); + +impl RegisterTile { + /// 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 + } +} + +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 + } +} + +/// 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, +{ + pub(crate) storage: &'a [Storage], + pub(crate) layout: Layout, + 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) + } +} + +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, +{ + storage: *mut Storage, + storage_len: usize, + layout: Layout, + element: PhantomData, + 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) + } +} + +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 using its absolute index in the original tensor. + /// + /// Use this for a short final tile. 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_linear(&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_div::(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_div::(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_linear`]. +#[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 }; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Clone, Copy)] + struct TestElement; + + impl TensorElement for TestElement {} + + #[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)); + } + + #[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()); + } + + #[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]); + } + + #[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/cuteir/cute-rs/src/tile.rs b/cuteir/cute-rs/src/tile.rs new file mode 100644 index 0000000000..081f8be0dd --- /dev/null +++ b/cuteir/cute-rs/src/tile.rs @@ -0,0 +1,104 @@ +/* + * 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`. +/// +/// ```text +/// x = 3 * 4 = 12 +/// assume_div::<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_div(x: usize) -> usize { + let _ = x; + unreachable!("cute-rs `assume_div` 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 diff --git a/cuteir/cute-rs/src/tiled_copy.rs b/cuteir/cute-rs/src/tiled_copy.rs new file mode 100644 index 0000000000..856694d3bd --- /dev/null +++ b/cuteir/cute-rs/src/tiled_copy.rs @@ -0,0 +1,225 @@ +/* + * 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::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. +//! These wrappers add no runtime storage beyond the source and destination +//! views. The compiler inlines them down to [`crate::copy_g2s`]. + +use core::marker::PhantomData; + +use crate::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 = ()> { + carrier: GmemMatrix, + 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 = ()> { + carrier: SmemTile, + 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::*; + + #[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. The destination layout comes from the +/// [`SharedTensor`] passed to [`Self::copy`]. +pub struct TiledCopy { + 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, + } + } + + /// Let the whole block copy one global tile into shared memory. + /// + /// `tile_coord` selects the source tile. `thread_idx` selects this thread's + /// assigned part of that tile. + /// + /// # 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, + thread_idx: u32, + ) { + unsafe { + copy_g2s::( + source.carrier(), + tile_coord, + destination.carrier_mut(), + thread_idx, + ) + } + } +} + +impl Default + for TiledCopy +{ + #[inline(always)] + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tiled_copy_is_static_and_typed_views_preserve_carrier_size() { + type Copy = TiledCopy<(), (), (), ()>; + assert_eq!(core::mem::size_of::(), 0); + #[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::>() + ); + } +} diff --git a/cuteir/cute-rs/src/tma.rs b/cuteir/cute-rs/src/tma.rs new file mode 100644 index 0000000000..49d248700d --- /dev/null +++ b/cuteir/cute-rs/src/tma.rs @@ -0,0 +1,528 @@ +/* + * 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: let desc: TmaDesc = make_tma_desc_2d(...)?; +//! kernel: fn kernel(desc: *const TmaDesc, ...) +//! └── same SmemL is required +//! ``` +//! +//! 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`. +#[repr(C, align(64))] +pub struct TmaDesc { + /// 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)>, +} + +/// 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. +/// +/// The default is [`Self::None`] to preserve the original behavior of +/// `make_tma_desc_2d`. CUTLASS and CuTeDSL commonly request [`Self::B128`] +/// 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 `make_tma_desc_2d_with_options`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TmaEncodeOptions { + /// L2 fetch size requested for this descriptor. + pub l2_promotion: TmaL2Promotion, +} + +impl TmaEncodeOptions { + /// Keep the original policy: do not request a larger L2 fetch. + pub const DEFAULT: Self = Self { + l2_promotion: TmaL2Promotion::None, + }; + + /// Match the usual CUTLASS/CuTeDSL input policy: fetch 128-byte 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::SWIZZLE_B, SmemL::SWIZZLE_M + log2, SmemL::SWIZZLE_S) +} + +/// 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), + } +} + +/// Number of bytes in one `SmemL` tile of `T`. +/// +/// 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 = tile_bytes::(); +/// ``` +#[inline(always)] +pub const fn tile_bytes() -> u32 { + (SmemL::ROWS * SmemL::COLS) 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::ROWS +/// col = c * SmemL::COLS +/// ``` +/// +/// The caller must initialize the `mbarrier`, attach [`tile_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 `make_tma_desc_2d` 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 TmaDesc, + coord: (usize, usize), + dst_smem_tile: &mut crate::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::ROWS, c * SmemL::COLS)`. +/// +/// 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::TmaStorePipeline`] +/// 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 TmaDesc, + coord: (usize, usize), + src_smem_tile: crate::SmemTile, +) { + let _ = (desc, coord, src_smem_tile); + unreachable!("cute-rs `copy_tma_s2g_2d` executed outside device compilation") +} + +#[cfg(feature = "host")] +mod host { + use super::{ + ReifySmem2D, TmaDesc, 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. + /// + /// 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 make_tma_desc_2d( + global_base: *mut core::ffi::c_void, + rows: u64, + cols: u64, + ld_elements: u64, + ) -> Result, alloc::string::String> { + make_tma_desc_2d_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 [`make_tma_desc_2d`]. 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 make_tma_desc_2d_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::ROWS <= 0 || SmemL::COLS <= 0 { + return Err(format!( + "TMA tile extents must be positive, got {} x {}", + SmemL::ROWS, + SmemL::COLS + )); + } + if SmemL::ROW_STRIDE != SmemL::COLS || SmemL::COL_STRIDE != 1 { + return Err(format!( + "TMA v0 requires a row-major shared tile, got strides ({}, {})", + SmemL::ROW_STRIDE, + SmemL::COL_STRIDE + )); + } + 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::SWIZZLE_B, + SmemL::SWIZZLE_M, + SmemL::SWIZZLE_S + )); + }; + if let Some(span_bytes) = tma_swizzle_span_bytes(mode) { + let pitch_bytes = SmemL::COLS + .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::COLS as u32, SmemL::ROWS 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() }; + const _: () = assert!(core::mem::size_of::() == 128); + Ok(TmaDesc { + bytes: unsafe { core::mem::transmute::(map) }, + marker: PhantomData, + }) + } +} + +#[cfg(feature = "host")] +pub use host::{TmaElement, make_tma_desc_2d, make_tma_desc_2d_with_options}; + +#[cfg(test)] +mod tests { + use super::*; + + #[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); + } + + #[test] + fn element_recast_reaches_the_byte_modes() { + use crate::markers::{Composed, ReifySmem2D, RowMajor, Swizzle}; + // For f16, element pattern S<3,3,3> becomes byte pattern S<3,4,3>, + // which is TMA B128. + type Smem = Composed, 0, RowMajor<8, 64>>; + assert_eq!(Smem::SWIZZLE_B, 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); + } + + #[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)); + } + + #[test] + fn flattened_mxfp4_ab_tile_is_8192_bytes_with_b64_swizzle() { + use crate::markers::{Composed, RowMajor, Swizzle}; + + // Two FP4 values share each byte: 128 K values become 64 bytes. + // Byte pattern S<2,4,3> is TMA B64. + type ByteSmem = Composed, 0, RowMajor<128, 64>>; + assert_eq!(tile_bytes::(), 8_192); + assert_eq!( + tma_swizzle_mode_for::(1), + Some(TmaSwizzleMode::B64) + ); + } + + #[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. + type ScaleSmem = RowMajor<1, 256>; + assert_eq!(tile_bytes::(), 512); + assert_eq!( + tma_swizzle_mode_for::(2), + Some(TmaSwizzleMode::None) + ); + } + + #[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); + } + + #[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/cuteir/cute-rs/tests/compile_fail/immutable_store.rs b/cuteir/cute-rs/tests/compile_fail/immutable_store.rs new file mode 100644 index 0000000000..8e68c11056 --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/immutable_store.rs @@ -0,0 +1,15 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![feature(f16)] + +use cute_rs::Tensor; + +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/cuteir/cute-rs/tests/compile_fail/immutable_store.stderr b/cuteir/cute-rs/tests/compile_fail/immutable_store.stderr new file mode 100644 index 0000000000..f1ee65d854 --- /dev/null +++ b/cuteir/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:14:19 + | +14 | unsafe { tile.store(values) }; + | ^^^^^ + | +help: there is a method `storage` with a similar name, but with different arguments + --> src/tensor.rs + | + | pub const fn storage(&self) -> &'a [Storage] { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/cuteir/cute-rs/tests/compile_fail/invalid_storage.rs b/cuteir/cute-rs/tests/compile_fail/invalid_storage.rs new file mode 100644 index 0000000000..9fb201701a --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/invalid_storage.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::{E2M1, KMajor, Mkl, Tensor}; + +fn main() { + let wrong_carrier = [0.0f32; 32]; + let _ = Tensor::::from_storage( + &wrong_carrier, + KMajor::::new(1, 64), + ); +} diff --git a/cuteir/cute-rs/tests/compile_fail/invalid_storage.stderr b/cuteir/cute-rs/tests/compile_fail/invalid_storage.stderr new file mode 100644 index 0000000000..700a895fc6 --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/invalid_storage.stderr @@ -0,0 +1,35 @@ +error[E0277]: the trait bound `E2M1: TensorElement` is not satisfied + --> tests/compile_fail/invalid_storage.rs:10:13 + | +10 | let _ = Tensor::::from_storage( + | ^^^^^^^^^^^^^^^^^^^^^^ the trait `TensorElement` is not implemented for `E2M1` + | +help: the trait `TensorElement` is not implemented for `E2M1` + but trait `TensorElement` is implemented for it + --> src/numeric.rs + | + | impl TensorElement for E2M1 {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + = help: for that trait implementation, expected `u8`, found `f32` +note: required by a bound in `Tensor` + --> src/tensor.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<'_, E2M1, _, f32>`, but its trait bounds were not satisfied + --> tests/compile_fail/invalid_storage.rs:10:37 + | +10 | let _ = Tensor::::from_storage( + | ^^^^^^^^^^^^ associated function or constant cannot be called on `Tensor<'_, E2M1, _, f32>` due to unsatisfied trait bounds + | + ::: src/numeric.rs + | + | pub struct E2M1; + | --------------- doesn't satisfy `E2M1: TensorElement` + | + = note: the following trait bounds were not satisfied: + `E2M1: TensorElement` diff --git a/cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.rs b/cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.rs new file mode 100644 index 0000000000..592db4ddc3 --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.rs @@ -0,0 +1,16 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![feature(f16)] + +use core::ptr; + +use cute_rs::{Mxf4E2M1, Mxfp4TiledMma, Nkl, SharedTensor}; + +fn main() { + let shared_b = unsafe { + SharedTensor::::from_raw_parts(ptr::null_mut(), 0) + }; + let mma = Mxfp4TiledMma::<()>::get_slice(0); + let _ = unsafe { mma.load_a(&shared_b, (0, 0)) }; +} diff --git a/cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.stderr b/cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.stderr new file mode 100644 index 0000000000..bbec330d8f --- /dev/null +++ b/cuteir/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:15:33 + | +15 | 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/block_scaled_mma.rs + | + | pub unsafe fn load_a( + | ^^^^^^ diff --git a/cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.rs b/cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.rs new file mode 100644 index 0000000000..d3e1f56456 --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.rs @@ -0,0 +1,24 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![feature(f16)] + +use cute_rs::{AccC, FragA, FragB, Mkl, Mxfp4BlockScaledTensor, Mxfp4TiledMma, Nkl}; + +fn main() { + let values = [f16::from_bits(0); 32]; + let scales = [127u8; 2]; + let matrix = Mxfp4BlockScaledTensor::::from_slices(&values, &scales, 1, 64); + let vector = Mxfp4BlockScaledTensor::::from_slices(&values, &scales, 1, 64); + let mma = Mxfp4TiledMma::<()>::get_slice(0); + + let _ = unsafe { + mma.gemm( + FragA([0; 4]), + vector.load_scale_pair(0, 0), + FragB([0; 2]), + matrix.load_scale_pair(0, 0), + AccC::ZERO, + ) + }; +} diff --git a/cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.stderr b/cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.stderr new file mode 100644 index 0000000000..f457c4425f --- /dev/null +++ b/cuteir/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:16:13 + | +16 | mma.gemm( + | ^^^^ +17 | FragA([0; 4]), +18 | vector.load_scale_pair(0, 0), + | ---------------------------- expected `MmaScalePair`, found `MmaScalePair` +19 | FragB([0; 2]), +20 | matrix.load_scale_pair(0, 0), + | ---------------------------- expected `MmaScalePair`, found `MmaScalePair` + | +note: method defined here + --> src/block_scaled_mma.rs + | + | pub unsafe fn gemm( + | ^^^^ +help: swap these arguments + | +16 ~ mma.gemm( +17 + FragA([0; 4]), +18 + matrix.load_scale_pair(0, 0), +19 + FragB([0; 2]), +20 + vector.load_scale_pair(0, 0), +21 + AccC::ZERO, +22 + ) + | diff --git a/cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.rs b/cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.rs new file mode 100644 index 0000000000..7f37992682 --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.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::{Contiguous1D, TensorMut}; + +fn require_clone() {} + +fn main() { + require_clone::>(); +} diff --git a/cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.stderr b/cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.stderr new file mode 100644 index 0000000000..00586a86c4 --- /dev/null +++ b/cuteir/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:11:21 + | +11 | 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:8:21 + | + 8 | fn require_clone() {} + | ^^^^^ required by this bound in `require_clone` diff --git a/cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.rs b/cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.rs new file mode 100644 index 0000000000..cee2d69c45 --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.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::{Contiguous1D, TensorMut}; + +fn require_copy() {} + +fn main() { + require_copy::>(); +} diff --git a/cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.stderr b/cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.stderr new file mode 100644 index 0000000000..38f037f791 --- /dev/null +++ b/cuteir/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:11:20 + | +11 | 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:8:20 + | + 8 | fn require_copy() {} + | ^^^^ required by this bound in `require_copy` diff --git a/cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.rs b/cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.rs new file mode 100644 index 0000000000..c1dad21f29 --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.rs @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![feature(f16)] + +use core::ptr; + +use cute_rs::{CpAsync, Mkl, Mxfp4BlockScaledTensor, SharedTensor, TiledCopy}; + +fn main() { + let values = [f16::from_bits(0); 32]; + let scales = [127u8; 2]; + let matrix = Mxfp4BlockScaledTensor::::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.copy(&source, (0, 0), &mut wrong_encoding, 0) }; +} diff --git a/cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.stderr b/cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.stderr new file mode 100644 index 0000000000..e04950d934 --- /dev/null +++ b/cuteir/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:20:41 + | +20 | unsafe { copy.copy(&source, (0, 0), &mut wrong_encoding, 0) }; + | ---- ^^^^^^^^^^^^^^^^^^^ 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/tiled_copy.rs + | + | pub unsafe fn copy< + | ^^^^ diff --git a/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.rs b/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.rs new file mode 100644 index 0000000000..523c3a5544 --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.rs @@ -0,0 +1,23 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![feature(f16)] + +use core::ptr; + +use cute_rs::{ + CpAsync, Mkl, Mxf4E2M1, Mxfp4BlockScaledTensor, Nkl, SharedTensor, TiledCopy, +}; + +fn main() { + let values = [f16::from_bits(0); 32]; + let scales = [127u8; 2]; + let matrix = Mxfp4BlockScaledTensor::::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.copy(&source, (0, 0), &mut shared_b, 0) }; +} diff --git a/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr b/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr new file mode 100644 index 0000000000..eadfa5e577 --- /dev/null +++ b/cuteir/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:22:41 + | +22 | unsafe { copy.copy(&source, (0, 0), &mut shared_b, 0) }; + | ---- ^^^^^^^^^^^^^ expected `&mut SharedTensor`, found `&mut SharedTensor<..., f16, (), ...>` + | | + | arguments to this method are incorrect + | + = note: expected mutable reference `&mut SharedTensor` + found mutable reference `&mut SharedTensor` +note: method defined here + --> src/tiled_copy.rs + | + | pub unsafe fn copy< + | ^^^^ diff --git a/cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.rs b/cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.rs new file mode 100644 index 0000000000..7e6c41b461 --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.rs @@ -0,0 +1,22 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![feature(f16)] + +use core::{marker::PhantomData, ptr}; + +use cute_rs::{RowMajor, SmemTile, TmaDesc, copy_tma_s2g_2d}; + +type DescriptorLayout = RowMajor<8, 8>; +type SourceLayout = RowMajor<4, 16>; + +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/cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.stderr b/cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.stderr new file mode 100644 index 0000000000..4ab478a8ff --- /dev/null +++ b/cuteir/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:21:50 + | +21 | 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/tma.rs + | + | pub unsafe fn copy_tma_s2g_2d( + | ^^^^^^^^^^^^^^^ diff --git a/cuteir/cute-rs/tests/type_safety.rs b/cuteir/cute-rs/tests/type_safety.rs new file mode 100644 index 0000000000..4e233bc99b --- /dev/null +++ b/cuteir/cute-rs/tests/type_safety.rs @@ -0,0 +1,18 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#[test] +fn tensor_and_block_scaled_types_fail_closed() { + let tests = trybuild::TestCases::new(); + 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/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"); +} diff --git a/cuteir/dialect-cute/Cargo.toml b/cuteir/dialect-cute/Cargo.toml new file mode 100644 index 0000000000..3358bc6ca9 --- /dev/null +++ b/cuteir/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/cuteir/dialect-cute/src/attributes.rs b/cuteir/dialect-cute/src/attributes.rs new file mode 100644 index 0000000000..dee1dc4f31 --- /dev/null +++ b/cuteir/dialect-cute/src/attributes.rs @@ -0,0 +1,1088 @@ +/* + * 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. + +use core::fmt; + +use combine::Parser; +use cute_layout::{ComposedLayout, Layout, OffsetUnit, ParseLayoutError, Swizzle}; +use pliron::common_traits::Verify; +use pliron::context::Context; +use pliron::impl_printable_for_display; +use pliron::irfmt::parsers::quoted_string_parser; +use pliron::parsable::{Parsable, ParseResult, StateStream}; +use pliron::result::Error; +use pliron::verify_err_noloc; +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(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(pub u32); + +/// 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 { + Producer, + 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 { + pub role: CutePipelineRoleAttr, + 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(()) + } +} + +/// 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 { + pub m_tiles: u64, + pub n_tiles: u64, + 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(()) + } +} + +/// 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 { + 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 { + ReadOnly, + 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 { + Plain, + E2M1, + UE8M0, +} + +/// 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 { + Generic, + Mkl, + 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 { + Contiguous1D, + Zipped1D(u64), + Tile1D(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(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(()) + } +} + +/// 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 { + Full, + Row, + KTile(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(()) + } +} + +/// 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(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(()) + } +} + +/// 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 +/// ``` +#[pliron_attr(name = "cute.matrix_role", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum CuteMatrixRoleAttr { + A, + B, +} + +/// A complete swizzled layout, stored as one typed attribute. +/// +/// The quoted payload is intentionally simple and round-trippable: +/// +/// ```text +/// "3,4,3;0;elements;(8,32):(32,1)" +/// B M S off unit inner layout +/// ``` +#[pliron_attr(name = "cute.composed_layout", verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteComposedLayoutAttr(pub ComposedLayout); + +impl fmt::Display for CuteComposedLayoutAttr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let swizzle = self.0.outer(); + write!( + f, + "\"{},{},{};{};{};{}\"", + swizzle.bits, + swizzle.base, + swizzle.shift, + self.0.offset(), + self.0.unit(), + self.0.inner() + ) + } +} + +impl_printable_for_display!(CuteComposedLayoutAttr); + +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 inner = 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 outer = Swizzle::try_new(bits, base, shift).ok_or(ParseLayoutError)?; + ComposedLayout::new(outer, offset, inner, 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() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use cute_layout::IntTuple; + + #[test] + fn composed_layout_text_round_trips_every_addressing_piece() { + let inner = 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, inner, 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); + } + + #[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()); + } + + #[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.outer(), Swizzle::new(2, 0, -3)); + assert_eq!(attr.to_string(), "\"2,0,-3;0;bytes;16:4\""); + } + + #[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()); + } +} + +/// 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(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() + } +} + +/// 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 { + 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 { + ScaleStage, + ScaleK64, + A, + 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 { + pub m: u64, + pub n: u64, + pub k: u64, + pub a_format: CuteTensorFormatAttr, + pub b_format: CuteTensorFormatAttr, + pub scale_format: CuteTensorFormatAttr, + pub values_per_scale: u64, + pub accumulator: CuteMmaAccumulatorAttr, + pub threads: u32, + pub a_registers: u32, + pub b_registers: u32, + pub accumulator_registers: u32, +} + +impl CuteMmaAtomAttr { + /// SM120's packed MXFP4 `m16n8k64` atom. + #[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, + } + } +} + +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 + || self.scale_format != CuteTensorFormatAttr::UE8M0 + { + return verify_err_noloc!("cute.mma_atom v0 needs E2M1 A/B values with UE8M0 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 { + pub atom: CuteMmaAtomAttr, + pub cta_m: u64, + pub cta_n: u64, + pub cta_k: u64, + pub warp_m: u32, + pub warp_n: u32, + pub b_load_group: u32, + pub shared_layout: CuteComposedLayoutAttr, + pub m_ownership: CuteLayoutAttr, + 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"), + ), + } + } + + /// 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)) + } +} + +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.inner().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.leaves().first().copied() != Some(i64::from(self.warp_m)) + || self.n_ownership.0.shape.leaves().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(()) + } +} + +/// 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 { + pub stages: u32, +} + +impl CuteTmaStorePipelineAttr { + #[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 { + pub barrier_id: u32, + pub first_warp: u32, + pub warp_count: u32, + pub cta_warps: u32, + pub lanes_per_warp: u32, +} + +impl CuteCountedCtaBarrierAttr { + #[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(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 { + pub tiled_mma: CuteTiledMmaPlanAttr, + pub half_layout: CuteComposedLayoutAttr, + pub halves: u32, + 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 inner: Layout = "(128,64):(64,1)" + .parse() + .expect("fixed epilogue half layout is valid"); + let half_layout = + ComposedLayout::new(Swizzle::new(3, 3, 3), 0, inner, OffsetUnit::Elements) + .expect("fixed epilogue swizzle is valid"); + Self { + tiled_mma, + half_layout: CuteComposedLayoutAttr(half_layout), + halves: 2, + base_alignment: CuteAlignmentAttr(1024), + } + } + + #[must_use] + pub fn half_elements(&self) -> Option { + u64::try_from(self.half_layout.0.inner().checked_size()?).ok() + } + + #[must_use] + pub fn full_elements(&self) -> Option { + self.half_elements()?.checked_mul(u64::from(self.halves)) + } + + #[must_use] + pub fn half_bytes(&self, element_bytes: u64) -> Option { + self.half_elements()?.checked_mul(element_bytes) + } + + #[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_inner: 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_inner, 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.inner().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/cuteir/dialect-cute/src/epilogue_ops.rs b/cuteir/dialect-cute/src/epilogue_ops.rs new file mode 100644 index 0000000000..191cbf7991 --- /dev/null +++ b/cuteir/dialect-cute/src/epilogue_ops.rs @@ -0,0 +1,1182 @@ +/* + * 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::location::Located; +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}; + +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()) +} + +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()) +} + +fn epilogue_tile_of_type(ctx: &Context, ty: TypeHandle) -> Option { + ty.deref(ctx) + .downcast_ref::() + .cloned() +} + +fn tma_view_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +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 + }) +} + +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) +} + +fn current_half_layout() -> ComposedLayout { + let inner: Layout = "(128,64):(64,1)" + .parse() + .expect("fixed epilogue half layout is valid"); + ComposedLayout::new(Swizzle::new(3, 3, 3), 0, inner, OffsetUnit::Elements) + .expect("fixed epilogue half swizzle is valid") +} + +/// 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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn input_base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + #[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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn input_base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[must_use] + pub fn input_warp_id(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + #[must_use] + pub fn input_lane(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(2) + } + + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + #[must_use] + pub fn warp_id(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + #[must_use] + pub fn lane(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(2) + } + + #[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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[must_use] + pub fn warp_id(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + #[must_use] + pub fn lane(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(2) + } + + #[must_use] + pub fn accumulator(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(3) + } + + #[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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[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)) + } + + #[must_use] + pub fn barrier(&self, ctx: &Context) -> Option { + self.get_attr_epilogue_sync_barrier(ctx) + .map(|barrier| *barrier) + } + + #[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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn full_base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[must_use] + pub fn half_base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + #[must_use] + pub fn capacity(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + #[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)) + } + + #[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(()) + } +} + +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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[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 { + 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, + ), + } + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn source(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[must_use] + pub fn destination(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + #[must_use] + pub fn tile_row(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(2) + } + + #[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(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::attributes::{CuteEpiloguePlanAttr, CuteTiledMmaPlanAttr}; + use crate::gemm_tma_ops::{CuteTmaGmemViewOp, CuteTmaSmemViewOp}; + use crate::types::CuteEpilogueTileType; + use dialect_mir::ops::MirUndefOp; + use dialect_mir::types::{MirArrayType, MirStructType}; + use pliron::builtin::types::FP32Type; + + fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) + } + + fn data_layout() -> ComposedLayout { + let inner: Layout = "(128,32):(32,1)".parse().unwrap(); + ComposedLayout::new(Swizzle::new(2, 3, 3), 0, inner, OffsetUnit::Elements).unwrap() + } + + struct Fixture { + base: Value, + warp: Value, + lane: Value, + accumulator: Value, + capacity: Value, + row: Value, + column: Value, + descriptor: Value, + f16_ty: TypeHandle, + tile_ty: TypeHandle, + } + + 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, + }, + ) + } + + #[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) + ); + } + + #[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() + ); + } + + #[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()); + } + + #[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()); + } + + #[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/cuteir/dialect-cute/src/gemm_tma_ops.rs b/cuteir/dialect-cute/src/gemm_tma_ops.rs new file mode 100644 index 0000000000..5d2f3e8450 --- /dev/null +++ b/cuteir/dialect-cute/src/gemm_tma_ops.rs @@ -0,0 +1,783 @@ +/* + * 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::location::Located; +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}; + +fn tma_view_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +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, + } + }) + } +} + +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()) +} + +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 { + 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, + ) + } + + 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 { + #[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, + ), + } + } + + 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 { + 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, + ), + } + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn source(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[must_use] + pub fn destination(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + #[must_use] + pub fn tile_row(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(2) + } + + #[must_use] + pub fn tile_column(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(3) + } + + #[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(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::layout::{Layout, OffsetUnit, Swizzle}; + use dialect_mir::ops::MirUndefOp; + use pliron::builtin::types::Signedness; + + fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) + } + + fn result(ctx: &Context, op: Ptr) -> Value { + op.deref(ctx).get_result(0) + } + + fn byte_layout() -> ComposedLayout { + let inner: Layout = "(128,64):(64,1)".parse().unwrap(); + ComposedLayout::new(Swizzle::new(2, 4, 3), 0, inner, OffsetUnit::Elements).unwrap() + } + + fn scale_layout() -> ComposedLayout { + let inner: Layout = "(1,256):(256,1)".parse().unwrap(); + ComposedLayout::new(Swizzle::IDENTITY, 0, inner, OffsetUnit::Elements).unwrap() + } + + struct Inputs { + descriptor: Value, + byte_base: Value, + scale_base: Value, + size: Value, + barrier: Value, + u8_ty: TypeHandle, + u16_ty: TypeHandle, + } + + 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, + }, + ) + } + + 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) + } + + #[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()); + } + } + + #[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)); + } + + #[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") + ); + } + + #[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") + ); + } + + #[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") + ); + } + + #[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") + ); + } + + #[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/cuteir/dialect-cute/src/gemv_ops.rs b/cuteir/dialect-cute/src/gemv_ops.rs new file mode 100644 index 0000000000..44187b411f --- /dev/null +++ b/cuteir/dialect-cute/src/gemv_ops.rs @@ -0,0 +1,1095 @@ +/* + * 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::location::Located; +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; + +fn tensor_view_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +fn scaled_view_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +fn fragment_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +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()) +} + +fn is_f32(ctx: &Context, ty: TypeHandle) -> bool { + ty.deref(ctx).downcast_ref::().is_some() +} + +fn same_scaled_members(left: &CuteScaledViewType, right: &CuteScaledViewType) -> bool { + left.values == right.values && left.scales == right.scales && left.role == right.role +} + +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 +} + +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) +} + +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), + ) + } + + 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, + CuteTensorLayoutAttr::BlockScaleKMajor(_) + ) + ) { + return verify_err!( + op.loc(), + "cute.tensor_make_2d needs an E2M1 value or UE8M0 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 { + 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, + ), + } + } + + 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 { + 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, + ), + } + } + + 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 { + 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, + ), + } + } + + 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 { + 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 + } + + 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 { + 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, + ), + } + } + + 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(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use dialect_mir::ops::MirUndefOp; + use dialect_mir::types::MirFP16Type; + + fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) + } + + struct Inputs { + matrix_values: Value, + matrix_scales: Value, + vector_values: Value, + vector_scales: Value, + size: Value, + acc: Value, + } + + 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, + }, + ) + } + + fn result(ctx: &Context, op: Ptr) -> Value { + op.deref(ctx).get_result(0) + } + + 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) + } + + 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) + } + + #[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) + ); + } + + #[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") + ); + } + + #[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") + ); + } + + #[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>") + ); + } + + #[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") + ); + } + + #[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()); + } + + #[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/cuteir/dialect-cute/src/lib.rs b/cuteir/dialect-cute/src/lib.rs new file mode 100644 index 0000000000..07d1dd8857 --- /dev/null +++ b/cuteir/dialect-cute/src/lib.rs @@ -0,0 +1,142 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! CuTe dialect definition. +//! +//! Models the CuTe layout algebra as first-class pliron ops, types, and +//! attributes, coexisting in one module with dialect-mir control flow and +//! scalar operations. +//! +//! Layouts are stored as typed attributes backed by `cute-layout` (never as +//! strings); result types of algebra ops are computed with `cute-layout` at +//! IR-build time. The backend-neutral verifier checks the complete semantic +//! graph before any backend continuation begins. +//! +//! 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, and the shared-tensor/tiled-MMA compute story. These semantic +//! operations are later consumed by the selected backend continuation. + +pub mod attributes; +pub mod epilogue_ops; +pub mod gemm_tma_ops; +pub mod gemv_ops; +pub mod ops; +pub mod pipeline_ops; +pub mod scheduler_ops; +pub mod smem_mma_ops; +pub mod tensor_ops; +pub mod types; +pub mod verify; + +/// Re-export of the layout algebra so consumers (mir-importer's recognition +/// module and backend continuations) depend only on `dialect-cute` and reach +/// the algebra through it. +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; + +pub const CUTE_DIALECT_NAME: &str = "cute"; + +pub fn register(ctx: &mut Context) { + Dialect::register( + ctx, + &DialectName::try_new(CUTE_DIALECT_NAME).expect("valid dialect name"), + ); + + // The #[pliron_op]/#[pliron_attr] macros auto-register via Context::default; + // explicit registration is the idempotent house convention. + ops::CuteCopyOp::register(ctx); + ops::CuteAssumeDivOp::register(ctx); + ops::CuteCopyG2SOp::register(ctx); + ops::CuteLdmatrixOp::register(ctx); + ops::CuteTmaLoad2dOp::register(ctx); + ops::CuteTmaStore2dOp::register(ctx); + tensor_ops::CuteTensorMakeOp::register(ctx); + tensor_ops::CuteTensorZippedDivideOp::register(ctx); + tensor_ops::CuteTensorSliceOp::register(ctx); + tensor_ops::CuteTensorIsFullOp::register(ctx); + tensor_ops::CuteTensorBaseOp::register(ctx); + tensor_ops::CuteTensorLoadIntoOp::register(ctx); + tensor_ops::CuteTensorStoreFromOp::register(ctx); + tensor_ops::CuteTensorStoreElementAbsOp::register(ctx); + gemv_ops::CuteTensorMake2DOp::register(ctx); + gemv_ops::CuteScaledViewMakeOp::register(ctx); + gemv_ops::CuteScaledViewRowOp::register(ctx); + gemv_ops::CuteScaledViewKTileOp::register(ctx); + gemv_ops::CuteScaledViewLoadOp::register(ctx); + gemv_ops::CuteDotOp::register(ctx); + gemm_tma_ops::CuteTmaGmemViewOp::register(ctx); + gemm_tma_ops::CuteTmaSmemViewOp::register(ctx); + gemm_tma_ops::CuteTmaCopy2dOp::register(ctx); + epilogue_ops::CuteEpilogueSmemOverlayOp::register(ctx); + epilogue_ops::CuteEpilogueWarpSliceOp::register(ctx); + epilogue_ops::CuteEpilogueStoreFragmentOp::register(ctx); + epilogue_ops::CuteEpilogueSyncOp::register(ctx); + epilogue_ops::CuteEpilogueHalfOp::register(ctx); + epilogue_ops::CuteTmaStoreAcquireOp::register(ctx); + epilogue_ops::CuteTmaStoreCommitOp::register(ctx); + epilogue_ops::CuteTmaStoreTailOp::register(ctx); + epilogue_ops::CuteTmaStore2dSemanticOp::register(ctx); + scheduler_ops::CuteSchedulerNew1dOp::register(ctx); + scheduler_ops::CuteSchedulerHasWorkOp::register(ctx); + scheduler_ops::CuteSchedulerCurrentOp::register(ctx); + scheduler_ops::CuteWorkTileCoordinatesOp::register(ctx); + scheduler_ops::CuteSchedulerAdvanceOp::register(ctx); + pipeline_ops::CuteTmaLoadPipelineMakeOp::register(ctx); + pipeline_ops::CuteTmaLoadPipelineInitOp::register(ctx); + pipeline_ops::CutePipelineStateNewOp::register(ctx); + pipeline_ops::CutePipelineStateSlotOp::register(ctx); + pipeline_ops::CutePipelineStateAdvanceOp::register(ctx); + pipeline_ops::CutePipelineProducerAcquireOp::register(ctx); + pipeline_ops::CutePipelineProducerExpectTxOp::register(ctx); + pipeline_ops::CutePipelineConsumerWaitOp::register(ctx); + pipeline_ops::CutePipelineConsumerReleaseOp::register(ctx); + pipeline_ops::CutePipelineProducerTailOp::register(ctx); + smem_mma_ops::CuteSmemTensorOverlayOp::register(ctx); + smem_mma_ops::CuteTiledMmaSliceOp::register(ctx); + smem_mma_ops::CuteFragmentFillOp::register(ctx); + smem_mma_ops::CuteMmaLoadScalesOp::register(ctx); + smem_mma_ops::CuteFragmentSliceKOp::register(ctx); + smem_mma_ops::CuteMmaLoadAOp::register(ctx); + smem_mma_ops::CuteMmaPartitionBOp::register(ctx); + smem_mma_ops::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/cuteir/dialect-cute/src/ops.rs b/cuteir/dialect-cute/src/ops.rs new file mode 100644 index 0000000000..60b193bba2 --- /dev/null +++ b/cuteir/dialect-cute/src/ops.rs @@ -0,0 +1,1509 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! cute dialect ops. +//! +//! The op set is intentionally small. `cute.copy` and `cute.copy_g2s` are +//! zero-result memory ops; `cute.assume_div` is an identity value that +//! preserves one compiler fact. + +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, OneOpdInterface, OneResultInterface}, + types::IntegerType, + }, + common_traits::Verify, + context::{Context, Ptr}, + location::Located, + 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, + CuteMatrixRoleAttr, +}; + +/// 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 { + #[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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +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(()) + } +} + +/// Preserve the promise `value % divisor == 0` in SSA form. +/// +/// ```text +/// Rust cute IR lowering +/// ───── ─────── ──────── +/// y = assume_div::<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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + 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(()) + } +} + +/// 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(()) + } +} + +/// 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`. +#[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 { + 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 + } + + 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.inner().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(()) + } +} + +/// 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 { + 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 + } + + 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 { + 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 + } + + 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(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use cute_layout::{IntTuple, MAX_STATIC_VALIDATION_BYTES, OffsetUnit, Swizzle}; + use dialect_mir::ops::MirUndefOp; + use pliron::builtin::types::{FP32Type, Signedness}; + + fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) + } + + 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) + } + + #[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, + ) + } + + #[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, + ) + } + + 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}`" + ); + } + + 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) + } + + 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) + } + + #[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}" + ); + } + + #[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(); + } + + #[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}"); + } + + #[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}"); + } + + #[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"); + } + + #[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"); + } + + #[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", + ); + } + + #[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", + ); + } + + #[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"); + } + + #[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"); + } + + #[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"); + } + + #[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/cuteir/dialect-cute/src/pipeline_ops.rs b/cuteir/dialect-cute/src/pipeline_ops.rs new file mode 100644 index 0000000000..e817f3158d --- /dev/null +++ b/cuteir/dialect-cute/src/pipeline_ops.rs @@ -0,0 +1,1376 @@ +/* + * 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::location::Located; +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; + +const BARRIER_BYTES: u64 = 8; + +fn u32_type(ctx: &Context) -> TypeHandle { + IntegerType::get(ctx, 32, Signedness::Unsigned).into() +} + +fn u64_type(ctx: &Context) -> TypeHandle { + IntegerType::get(ctx, 64, Signedness::Unsigned).into() +} + +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()) +} + +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()) +} + +fn pipeline_type_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +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) +} + +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)) +} + +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)) +} + +fn barrier_pointer_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +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()) +} + +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) +} + +/// 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 + } + + 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 { + 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, + ), + } + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[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(()) + } +} + +/// 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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn slot(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + #[must_use] + pub fn phase(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + #[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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[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) + } + + #[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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn slot(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[must_use] + pub fn phase(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + #[must_use] + pub fn next_slot(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + #[must_use] + pub fn next_phase(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + #[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) + } +} + +macro_rules! pipeline_state_accessors { + ($attr_getter:ident) => { + #[must_use] + pub fn pipeline(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[must_use] + pub fn slot(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + #[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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + pipeline_state_accessors!(get_attr_pipeline_producer_acquire_state); + + #[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 { + 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 + } + + 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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + pipeline_state_accessors!(get_attr_pipeline_consumer_wait_state); + + #[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 { + 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 + } + + 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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + pipeline_state_accessors!(get_attr_pipeline_producer_tail_state); + + #[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(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use dialect_mir::ops::MirUndefOp; + + const STAGES: u64 = 3; + const CONSUMER_WARPS: u32 = 8; + const TRANSACTION_BYTES: u32 = 17_408; + + fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) + } + + 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) + } + + struct Setup { + ctx: Context, + base: Value, + u32_value: Value, + u64_value: Value, + } + + 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, + } + } + + #[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) + ); + } + + #[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()); + } + + #[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()); + } + + #[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") + ); + } + + #[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") + ); + } + + #[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/cuteir/dialect-cute/src/scheduler_ops.rs b/cuteir/dialect-cute/src/scheduler_ops.rs new file mode 100644 index 0000000000..b897dafc65 --- /dev/null +++ b/cuteir/dialect-cute/src/scheduler_ops.rs @@ -0,0 +1,598 @@ +/* + * 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 ──► current_tile ──► (linear, m, n, batch) +//! ▲ │ +//! └──────────── advance(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::location::Located; +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; + +fn u64_type(ctx: &Context) -> TypeHandle { + IntegerType::get(ctx, 64, Signedness::Unsigned).into() +} + +fn bool_type(ctx: &Context) -> TypeHandle { + IntegerType::get(ctx, 1, Signedness::Signless).into() +} + +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()) +} + +fn is_bool(ctx: &Context, value: Value) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 1) +} + +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 { + 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 + } + + 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 { + 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 + } + + 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 { + 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 + } + + 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 { + 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, + ), + } + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn work_tile(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[must_use] + pub fn linear(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + #[must_use] + pub fn m_tile(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + #[must_use] + pub fn n_tile(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(2) + } + + #[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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn current(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[must_use] + pub fn stride(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + #[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) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use dialect_mir::ops::MirUndefOp; + + fn result(ctx: &Context, op: Ptr, index: usize) -> Value { + op.deref(ctx).get_result(index) + } + + fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) + } + + #[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))); + } + + #[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") + ); + } + + #[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") + ); + } + + #[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()); + } + + #[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/cuteir/dialect-cute/src/smem_mma_ops.rs b/cuteir/dialect-cute/src/smem_mma_ops.rs new file mode 100644 index 0000000000..e1fc3fa1ef --- /dev/null +++ b/cuteir/dialect-cute/src/smem_mma_ops.rs @@ -0,0 +1,1571 @@ +/* + * 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::location::Located; +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}; + +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()) +} + +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()) +} + +fn is_f32(ctx: &Context, value: Value) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some() +} + +fn smem_tensor_of_type(ctx: &Context, ty: TypeHandle) -> Option { + ty.deref(ctx).downcast_ref::().cloned() +} + +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 + }) +} + +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) + }) +} + +fn canonical_scale_placement() -> ComposedLayout { + let inner: Layout = "(32,4):(4,1)" + .parse() + .expect("fixed SM120 scale placement is valid"); + ComposedLayout::new(Swizzle::IDENTITY, 0, inner, OffsetUnit::Elements) + .expect("identity scale placement is valid") +} + +fn canonical_data_placement() -> ComposedLayout { + let inner: Layout = "(128,32):(32,1)" + .parse() + .expect("fixed SM120 data placement is valid"); + ComposedLayout::new(Swizzle::new(2, 3, 3), 0, inner, OffsetUnit::Elements) + .expect("fixed SM120 data swizzle is valid") +} + +fn plan_of_attr( + plan: Option>, +) -> Option { + plan.map(|plan| plan.clone()) +} + +fn check_plan(plan: &CuteTiledMmaPlanAttr, ctx: &Context) -> Result<(), &'static str> { + plan.verify(ctx) + .map_err(|_| "has an invalid tiled-MMA plan")?; + if plan.atom != crate::attributes::CuteMmaAtomAttr::mxf4_m16n8k64() + || 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 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(()) +} + +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(()) +} + +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 != CuteTensorFormatAttr::UE8M0 + || 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 shared scale atom"); + } + Ok(()) +} + +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") + } +} + +/// 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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn input_base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[must_use] + pub fn input_capacity(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + #[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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn lane(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[must_use] + pub fn sliced_lane(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + #[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(()) + } +} + +/// 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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn fill(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + #[must_use] + pub fn fragment(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + #[must_use] + pub fn plan(&self, ctx: &Context) -> Option { + plan_of_attr(self.get_attr_fragment_fill_plan(ctx)) + } + + #[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(()) + } +} + +/// 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 { + #[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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[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)) + } + + #[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(()) + } +} + +/// 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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[must_use] + pub fn plan(&self, ctx: &Context) -> Option { + plan_of_attr(self.get_attr_fragment_slice_k_plan(ctx)) + } + + #[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(()) + } +} + +/// 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 { + #[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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[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)) + } + + #[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 { + 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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[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)) + } + + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + #[must_use] + pub fn capacity(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + #[must_use] + pub fn warp_n(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(2) + } + + #[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(()) + } +} + +/// 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 { + #[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 + } + + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + #[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)) + } + + #[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(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::attributes::{CuteTensorAccessAttr, CuteTensorAddressSpaceAttr}; + use crate::types::CuteTensorViewType; + use dialect_mir::ops::MirUndefOp; + use pliron::builtin::types::Signedness; + + fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) + } + + fn data_placement() -> ComposedLayout { + let inner: Layout = "(128,32):(32,1)".parse().unwrap(); + ComposedLayout::new(Swizzle::new(2, 3, 3), 0, inner, OffsetUnit::Elements).unwrap() + } + + struct SmemTypeFacts { + format: CuteTensorFormatAttr, + role: CuteTensorRoleAttr, + alignment: u64, + layout: CuteTensorLayoutAttr, + placement: ComposedLayout, + } + + 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() + } + + struct TestTypes { + u32_ty: TypeHandle, + u64_ty: TypeHandle, + f32_ty: TypeHandle, + a_ptr_ty: TypeHandle, + b_ptr_ty: TypeHandle, + scale_a_ptr_ty: TypeHandle, + scale_b_ptr_ty: TypeHandle, + a_view_ty: TypeHandle, + b_view_ty: TypeHandle, + scale_a_view_ty: TypeHandle, + scale_b_view_ty: TypeHandle, + a_carrier_ty: TypeHandle, + scale_carrier_ty: TypeHandle, + accumulator_ty: TypeHandle, + plan: CuteTiledMmaPlanAttr, + } + + 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, + }, + ) + } + + 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) + } + + #[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)) + ); + } + + #[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") + ); + } + + #[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()); + } + + #[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") + ); + } + + #[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/cuteir/dialect-cute/src/tensor_ops.rs b/cuteir/dialect-cute/src/tensor_ops.rs new file mode 100644 index 0000000000..9ac2290dd0 --- /dev/null +++ b/cuteir/dialect-cute/src/tensor_ops.rs @@ -0,0 +1,952 @@ +/* + * 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::location::Located; +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; + +fn view_of(ctx: &Context, value: Value) -> Option { + let ty = value.get_type(ctx); + ty.deref(ctx).downcast_ref::().cloned() +} + +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() +} + +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()) +} + +fn is_i1(ctx: &Context, ty: TypeHandle) -> bool { + ty.deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 1) +} + +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 +} + +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) +} + +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, + ), + } + } + + 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 { + 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 + } + + 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 { + 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, + ), + } + } + + 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 { + 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, + ), + } + } + + 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 { + 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, + ), + } + } + + 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 { + 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 + } + + 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 + } + + 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 { + 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, + ), + } + } + + 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(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use dialect_mir::ops::MirUndefOp; + use pliron::builtin::types::FP32Type; + + fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) + } + + 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) + } + + #[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() + ); + } + + #[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") + ); + } + + #[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/cuteir/dialect-cute/src/types.rs b/cuteir/dialect-cute/src/types.rs new file mode 100644 index 0000000000..53b589dfaf --- /dev/null +++ b/cuteir/dialect-cute/src/types.rs @@ -0,0 +1,1222 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Tensor-view types kept only while CuTe meaning is useful. +//! +//! 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. No backend may assign this ghost type an ABI or leave +//! it in a kernel signature. +//! +//! Backend-neutral verification reads these values after common MIR +//! preparation has exposed the semantic story: +//! +//! ```text +//! make -> zipped_divide -> slice -> load/store +//! ``` +//! +//! At that seam, a view cannot sit inside a pointer or aggregate, cross a +//! block edge, or appear in a function signature. Ordinary scalar carriers +//! may remain in closed compiler-owned local cells or cross CFG edges; the +//! semantic verifier follows those carriers before backend selection. + +use dialect_mir::types::{MirArrayType, MirFP16Type, MirStructType, MirTupleType}; +use pliron::builtin::types::{FP32Type, IntegerType}; +use pliron::common_traits::Verify; +use pliron::context::Context; +use pliron::derive::pliron_type; +use pliron::location::Location; +use pliron::result::Error; +use pliron::r#type::{Type, TypeHandle, TypedHandle}; +use pliron::verify_err; + +use crate::attributes::{ + CuteAlignmentAttr, CuteComposedLayoutAttr, CuteEpiloguePlanAttr, CuteScaledLayoutAttr, + CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, CuteTensorFormatAttr, CuteTensorLayoutAttr, + CuteTensorRoleAttr, CuteTileGridAttr, +}; + +/// Physical scalar registers inside one ordinary MIR aggregate. +/// +/// Semantic MMA operations keep their runtime values in the same Rust-shaped +/// aggregates used today. This small summary lets their verifiers check those +/// carriers without turning them into new loop-carried ghost values. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct MmaCarrierShape { + pub u32_registers: u64, + pub f32_registers: u64, +} + +impl MmaCarrierShape { + #[must_use] + pub const fn u32(registers: u64) -> Self { + Self { + u32_registers: registers, + f32_registers: 0, + } + } + + #[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) + } + + 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)?, + }) + } + + 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 fails closed with `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)?) + }) +} + +/// 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 { + pub logical: TypeHandle, + pub storage: TypeHandle, + pub format: CuteTensorFormatAttr, + pub role: CuteTensorRoleAttr, + pub space: CuteTensorAddressSpaceAttr, + pub access: CuteTensorAccessAttr, + pub alignment: CuteAlignmentAttr, + 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, + ) + } + + /// Size of one v0 storage element in bytes. + #[must_use] + pub fn storage_bytes(&self, ctx: &Context) -> Option { + let storage = self.storage.deref(ctx); + if storage.downcast_ref::().is_some() { + Some(2) + } else if storage.downcast_ref::().is_some() { + Some(4) + } else if let Some(integer) = storage.downcast_ref::() { + let width = u64::from(integer.width()); + (width > 0 && width % 8 == 0).then_some(width / 8) + } else { + None + } + } + + /// 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, + } + } +} + +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" + ); + } + let storage_ty = self.storage.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() + || logical_is_f32 + || plain_integer_is_supported; + if !plain_is_supported { + return verify_err!( + Location::Unknown, + "plain cute.tensor_view supports u8, u16, f16, and f32 carriers 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 => { + 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 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 cute.tensor_view must be read-only" + ); + } + if !matches!(self.role, CuteTensorRoleAttr::Mkl | CuteTensorRoleAttr::Nkl) { + return verify_err!( + Location::Unknown, + "UE8M0 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 cute.tensor_view must use BlockScaleKMajor" + ); + } + } + } + Ok(()) + } +} + +/// 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 { + pub tensor: TypeHandle, + 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.inner().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, + CuteTensorLayoutAttr::BlockScaleKMajor(_) + ) + ) { + return verify_err!( + Location::Unknown, + "cute.smem_tensor needs packed E2M1 values or UE8M0 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(()) + } +} + +/// 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 { + pub storage: TypeHandle, + 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) + } + + #[must_use] + pub fn half_elements(&self) -> Option { + self.plan.half_elements() + } + + #[must_use] + pub fn full_elements(&self) -> Option { + self.plan.full_elements() + } + + #[must_use] + pub fn half_bytes(&self) -> Option { + self.plan.half_bytes(2) + } + + #[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(()) + } +} + +/// 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 { + pub tensor: TypeHandle, + 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.inner().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(()) + } +} + +/// 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 { + pub values: TypeHandle, + pub scales: TypeHandle, + pub role: CuteTensorRoleAttr, + 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 scales.format != CuteTensorFormatAttr::UE8M0 + || scales.layout != CuteTensorLayoutAttr::BlockScaleKMajor(16) + { + return verify_err!( + Location::Unknown, + "cute.scaled_view scales must be a UE8M0 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 { + 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(()) + } +} + +/// 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 `WorkTile` struct reaches the +/// runtime ABI. +/// +/// ```text +/// current linear number + grid<16,16,1> +/// │ +/// ▼ +/// !cute.work_tile> +/// │ coordinates +/// ▼ +/// linear, m, n, batch +/// ``` +#[pliron_type(name = "cute.work_tile", format = "`<` $grid `>`")] +#[derive(Hash, PartialEq, Eq, Debug, Clone)] +pub struct CuteWorkTileType { + 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) + } +} + +/// 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 { + pub stages: u64, + pub consumer_warps: u32, + 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(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::layout::{ComposedLayout, Layout, OffsetUnit, Swizzle}; + use pliron::builtin::types::Signedness; + + #[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()); + } + + #[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") + ); + } + + #[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)); + } + + #[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()); + } + + #[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/cuteir/dialect-cute/src/verify.rs b/cuteir/dialect-cute/src/verify.rs new file mode 100644 index 0000000000..d9c2f4eb55 --- /dev/null +++ b/cuteir/dialect-cute/src/verify.rs @@ -0,0 +1,4537 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Backend-neutral verification for complete high-level CuTe stories. +//! +//! This module is deliberately analysis-only. It follows semantic SSA +//! provenance, checks closed use graphs and ordered protocols, and audits +//! compiler-only CuTe types without cloning or rewriting the input module. + +use std::collections::{HashMap, HashSet}; + +use dialect_mir::attributes::MirCastKindAttr; +use dialect_mir::ops::control_flow::MirCondBranchOp; +use dialect_mir::ops::function::MirFuncOp; +use dialect_mir::ops::{ + MirAddOp, MirAllocaOp, MirCallOp, MirCastOp, MirConstantOp, MirConstructStructOp, + MirConstructTupleOp, MirExtractFieldOp, MirFieldAddrOp, MirGotoOp, MirInsertFieldOp, MirLoadOp, + MirNotOp, MirRefOp, MirStorageDeadOp, MirStorageLiveOp, MirStoreOp, MirUndefOp, +}; +use dialect_mir::types::{ + MirArrayType, MirDisjointSliceType, MirEnumType, MirPtrType, MirSliceType, MirStructType, + MirTupleType, MirUnionType, address_space, +}; +use pliron::attribute::attr_cast; +use pliron::basic_block::BasicBlock; +use pliron::builtin::attr_interfaces::TypedAttrInterface; +use pliron::builtin::op_interfaces::BranchOpInterface; +use pliron::builtin::type_interfaces::FunctionTypeInterface; +use pliron::builtin::types::IntegerType; +use pliron::common_traits::Verify; +use pliron::context::{Context, Ptr}; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::{Op, OpId, op_cast}; +use pliron::operation::Operation; +use pliron::r#type::{TypeHandle, Typed, type_cast}; +use pliron::value::Value; + +use crate::attributes::{ + CuteEpilogueSyncPhaseAttr, CutePipelineRoleAttr, CutePipelineStateAttr, CuteTensorRoleAttr, + CuteTileGridAttr, CuteTiledMmaPlanAttr, +}; +use crate::epilogue_ops::{ + CuteEpilogueHalfOp, CuteEpilogueSmemOverlayOp, CuteEpilogueStoreFragmentOp, CuteEpilogueSyncOp, + CuteEpilogueWarpSliceOp, CuteTmaStore2dSemanticOp, CuteTmaStoreAcquireOp, CuteTmaStoreCommitOp, + CuteTmaStoreTailOp, +}; +use crate::gemm_tma_ops::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}; +use crate::gemv_ops::{ + CuteDotOp, CuteScaledViewKTileOp, CuteScaledViewLoadOp, CuteScaledViewMakeOp, + CuteScaledViewRowOp, CuteTensorMake2DOp, +}; +use crate::ops::{ + CuteAssumeDivOp, CuteCopyG2SOp, CuteCopyOp, CuteLdmatrixOp, CuteTmaLoad2dOp, CuteTmaStore2dOp, +}; +use crate::pipeline_ops::{ + CutePipelineConsumerReleaseOp, CutePipelineConsumerWaitOp, CutePipelineProducerAcquireOp, + CutePipelineProducerExpectTxOp, CutePipelineProducerTailOp, CutePipelineStateAdvanceOp, + CutePipelineStateNewOp, CutePipelineStateSlotOp, CuteTmaLoadPipelineInitOp, + CuteTmaLoadPipelineMakeOp, +}; +use crate::scheduler_ops::{ + CuteSchedulerAdvanceOp, CuteSchedulerCurrentOp, CuteSchedulerHasWorkOp, CuteSchedulerNew1dOp, + CuteWorkTileCoordinatesOp, +}; +use crate::smem_mma_ops::{ + CuteFragmentFillOp, CuteFragmentSliceKOp, CuteMmaLoadAOp, CuteMmaLoadScalesOp, + CuteMmaPartitionBOp, CuteSmemTensorOverlayOp, CuteTiledGemmOp, CuteTiledMmaSliceOp, +}; +use crate::tensor_ops::{ + CuteTensorBaseOp, CuteTensorIsFullOp, CuteTensorLoadIntoOp, CuteTensorMakeOp, + CuteTensorSliceOp, CuteTensorStoreElementAbsOp, CuteTensorStoreFromOp, + CuteTensorZippedDivideOp, +}; +use crate::types::{ + CuteEpilogueTileType, CuteFragmentType, CuteScaledViewType, CuteSmemTensorType, + CuteTensorViewType, CuteTmaLoadPipelineType, CuteTmaViewType, CuteWorkTileType, +}; + +/// 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 { + #[error("invalid CuTe semantic module: {0}")] + Invalid(String), + #[error("unsupported CuTe semantic operation: {0}")] + Unsupported(String), +} + +/// Result type shared by backend-neutral verification and backend adapters. +pub type VerifyResult = std::result::Result; + +type Result = VerifyResult; + +fn invalid(message: impl Into) -> VerifyError { + VerifyError::Invalid(message.into()) +} + +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); + } + } + } +} + +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(), + ] +} + +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(), + ] +} + +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(), + ] +} + +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(), + ] +} + +fn tma_semantic_ids() -> [OpId; 3] { + [ + CuteTmaGmemViewOp::get_opid_static(), + CuteTmaSmemViewOp::get_opid_static(), + CuteTmaCopy2dOp::get_opid_static(), + ] +} + +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(), + ] +} + +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(), + ] +} + +fn all_high_level_semantic_ids() -> Vec { + let mut ids = tensor_semantic_ids().to_vec(); + 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 +} + +fn verify_local_op(ctx: &Context, operation: Ptr) -> Result { + let opid = Operation::get_opid(operation, ctx); + macro_rules! verify_as { + ($ty:ty) => { + <$ty>::wrap(operation).verify(ctx) + }; + } + + let result = 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(()) +} + +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 { + 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()) +} + +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(()) +} + +#[derive(Clone, Copy)] +struct TensorViewState { + tile_size: Option, + tile_index: Option, +} + +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}`" + ))) +} + +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(()) +} + +#[derive(Clone, Copy)] +struct GemvTensorState { + rows: Value, + k: Value, + role: CuteTensorRoleAttr, +} + +#[derive(Clone, Copy)] +struct GemvSelectionState { + values: GemvTensorState, + scales: GemvTensorState, + batch: Option, + row: Option, + tile_index: Option, +} + +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, + }) +} + +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}`" + ))) +} + +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(()) +} + +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(()) +} + +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(()) +} + +/// 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 +} + +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() +} + +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 +} + +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. +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)) +} + +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)) +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct CarrierKey { + value: Value, + /// Empty means the value itself. Non-empty paths describe a field inside + /// an ordinary aggregate before a matching `mir.extract_field` exposes it. + path: Vec, +} + +#[derive(Clone)] +struct OrdinaryFacts { + facts: HashMap>, +} + +impl Default for OrdinaryFacts { + fn default() -> Self { + Self { + facts: HashMap::new(), + } + } +} + +impl OrdinaryFacts +where + F: Clone + Eq + std::hash::Hash, +{ + fn insert_path(&mut self, value: Value, path: Vec, fact: F) -> bool { + self.facts + .entry(CarrierKey { value, path }) + .or_default() + .insert(fact) + } + + fn insert(&mut self, value: Value, fact: F) -> bool { + self.insert_path(value, Vec::new(), fact) + } + + fn at_path(&self, value: Value, path: &[u32]) -> Option<&HashSet> { + self.facts.get(&CarrierKey { + value, + path: path.to_vec(), + }) + } + + fn at(&self, value: Value) -> Option<&HashSet> { + self.at_path(value, &[]) + } + + fn entries(&self, value: Value) -> Vec<(Vec, F)> { + self.facts + .iter() + .filter(|(key, _)| key.value == value) + .flat_map(|(key, facts)| facts.iter().cloned().map(|fact| (key.path.clone(), fact))) + .collect() + } + + 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 + } +} + +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) +} + +fn is_local_cell_pointer_alias_cast(ctx: &Context, operation: Ptr) -> bool { + is_provenance_preserving_pointer_cast(ctx, operation, false) +} + +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. +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. +fn local_cell_is_closed(ctx: &Context, root: Value) -> bool { + let mut worklist = vec![root]; + let mut visited = HashSet::new(); + while let Some(pointer) = worklist.pop() { + if !visited.insert(pointer) { + continue; + } + for r#use in pointer.uses(ctx) { + let user = r#use.user_op(); + let opid = Operation::get_opid(user, ctx); + let operation = user.deref(ctx); + if r#use.find_index(ctx) != 0 { + return false; + } + let forwards_pointer = operation.get_num_results() == 1 + && ((opid == MirCastOp::get_opid_static() + && is_local_cell_pointer_alias_cast(ctx, user)) + || (opid == MirFieldAddrOp::get_opid_static() + && MirFieldAddrOp::new(user) + .get_attr_field_index(ctx) + .is_some())); + if forwards_pointer { + worklist.push(operation.get_result(0)); + } else if opid == MirLoadOp::get_opid_static() { + if MirLoadOp::new(user).is_volatile(ctx) { + return false; + } + } else if opid == MirStoreOp::get_opid_static() { + if MirStoreOp::new(user).is_volatile(ctx) { + return false; + } + } else { + return false; + } + } + } + true +} + +/// Propagate semantic facts through the exact ordinary SSA adapters that can +/// remain after common CuTe preparation. This is deliberately a monotone +/// fixed-point transfer: loop block arguments become known once their entry +/// fact reaches them, then semantic advance operations can feed the same fact +/// around the backedge on a later iteration. +fn propagate_ordinary_plumbing( + ctx: &Context, + operations: &[Ptr], + facts: &mut OrdinaryFacts, +) -> bool +where + F: Clone + Eq + std::hash::Hash, +{ + let mut changed = false; + let mut blocks = HashSet::new(); + for operation in operations { + let operation_ref = operation.deref(ctx); + if let Some(block) = operation_ref.get_parent_block() { + blocks.insert(block); + } + let opid = Operation::get_opid(*operation, ctx); + if opid == MirCastOp::get_opid_static() + && is_semantic_pointer_carrier_cast(ctx, *operation) + && operation_ref.get_num_operands() == 1 + && operation_ref.get_num_results() == 1 + { + changed |= facts.copy_entries( + operation_ref.get_operand(0), + operation_ref.get_result(0), + &[], + &[], + ); + } else if (opid == MirConstructStructOp::get_opid_static() + || opid == MirConstructTupleOp::get_opid_static()) + && operation_ref.get_num_results() == 1 + { + let result = operation_ref.get_result(0); + for (index, operand) in operation_ref.operands().enumerate() { + changed |= facts.copy_entries(operand, result, &[], &[index as u32]); + } + } else if opid == MirInsertFieldOp::get_opid_static() + && operation_ref.get_num_operands() == 2 + && operation_ref.get_num_results() == 1 + { + let Some(index) = MirInsertFieldOp::new(*operation).get_attr_insert_index(ctx) else { + continue; + }; + let base = operation_ref.get_operand(0); + let inserted = operation_ref.get_operand(1); + let result = operation_ref.get_result(0); + for (path, fact) in facts.entries(base) { + if path.first().copied() != Some(index.0) { + changed |= facts.insert_path(result, path, fact); + } + } + changed |= facts.copy_entries(inserted, result, &[], &[index.0]); + } else if opid == MirExtractFieldOp::get_opid_static() + && operation_ref.get_num_operands() == 1 + && operation_ref.get_num_results() == 1 + { + let Some(index) = MirExtractFieldOp::new(*operation).get_attr_index(ctx) else { + continue; + }; + changed |= facts.copy_entries( + operation_ref.get_operand(0), + operation_ref.get_result(0), + &[index.0], + &[], + ); + } else if opid == MirStoreOp::get_opid_static() && operation_ref.get_num_operands() == 2 { + let store = MirStoreOp::new(*operation); + if !store.is_volatile(ctx) + && let Some((root, path)) = local_cell(ctx, store.address_opd(ctx)) + && local_cell_is_closed(ctx, root) + { + changed |= facts.copy_entries(store.value_opd(ctx), root, &[], &path); + } + } else if opid == MirLoadOp::get_opid_static() + && operation_ref.get_num_operands() == 1 + && operation_ref.get_num_results() == 1 + { + let load = MirLoadOp::new(*operation); + if !load.is_volatile(ctx) + && let Some((root, path)) = local_cell(ctx, load.address_opd(ctx)) + && local_cell_is_closed(ctx, root) + { + changed |= facts.copy_entries(root, operation_ref.get_result(0), &path, &[]); + } + } + } + + for block in blocks { + let block_ref = block.deref(ctx); + for argument_index in 0..block_ref.get_num_arguments() { + let argument = block_ref.get_argument(argument_index); + for edge in block.uses(ctx) { + let terminator = edge.user_op(); + let Some(operands) = + branch_successor_operands(ctx, terminator, edge.find_index(ctx)) + else { + continue; + }; + if let Some(incoming) = operands.get(argument_index).copied() { + changed |= facts.copy_entries(incoming, argument, &[], &[]); + } + } + } + } + changed +} + +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(); + } + } +} + +fn seed_unknown_integer_carriers( + ctx: &Context, + operations: &[Ptr], + facts: &mut OrdinaryFacts, + width: u32, + unknown: impl Fn(Value) -> F, +) where + F: Clone + Eq + std::hash::Hash, +{ + let mut values = Vec::new(); + let mut blocks = HashSet::new(); + for operation in operations { + let operation_ref = operation.deref(ctx); + values.extend(operation_ref.results()); + if let Some(block) = operation_ref.get_parent_block() + && blocks.insert(block) + { + let block_ref = block.deref(ctx); + values.extend( + (0..block_ref.get_num_arguments()).map(|index| block_ref.get_argument(index)), + ); + } + } + + for value in values { + let mut paths = Vec::new(); + integer_leaf_paths(ctx, value.get_type(ctx), width, &mut Vec::new(), &mut paths); + for path in paths { + if facts.at_path(value, &path).is_none() { + facts.insert_path(value, path, unknown(value)); + } + } + } +} + +fn poison_missing_forwarding_sources( + ctx: &Context, + operations: &[Ptr], + facts: &mut OrdinaryFacts, + unknown: impl Fn(Value) -> F + Copy, +) -> bool +where + F: Clone + Eq + std::hash::Hash, +{ + let mut changed = false; + let mut blocks = HashSet::new(); + for operation in operations { + if let Some(block) = operation.deref(ctx).get_parent_block() { + blocks.insert(block); + } + if Operation::get_opid(*operation, ctx) != MirStoreOp::get_opid_static() { + continue; + } + let store = MirStoreOp::new(*operation); + let Some((root, cell_path)) = local_cell(ctx, store.address_opd(ctx)) else { + continue; + }; + if !local_cell_is_closed(ctx, root) { + continue; + } + for (path, _) in facts.entries(root) { + let Some(source_path) = path.strip_prefix(cell_path.as_slice()) else { + continue; + }; + if facts.at_path(store.value_opd(ctx), source_path).is_none() { + changed |= facts.insert_path( + store.value_opd(ctx), + source_path.to_vec(), + unknown(store.value_opd(ctx)), + ); + } + } + } + + for block in blocks { + let block_ref = block.deref(ctx); + for argument_index in 0..block_ref.get_num_arguments() { + let argument = block_ref.get_argument(argument_index); + let argument_entries = facts.entries(argument); + if argument_entries.is_empty() { + continue; + } + for edge in block.uses(ctx) { + let terminator = edge.user_op(); + let Some(operands) = + branch_successor_operands(ctx, terminator, edge.find_index(ctx)) + else { + continue; + }; + let Some(incoming) = operands.get(argument_index).copied() else { + continue; + }; + for (path, _) in &argument_entries { + if facts.at_path(incoming, path).is_none() { + changed |= facts.insert_path(incoming, path.clone(), unknown(incoming)); + } + } + } + } + } + changed +} + +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"))) +} + +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +enum SchedulerFact { + Current(Ptr), + Stride(Ptr), + Unknown(Value), +} + +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" + ))), + } +} + +fn build_scheduler_facts( + ctx: &Context, + operations: &[Ptr], +) -> OrdinaryFacts { + let mut facts = OrdinaryFacts::default(); + for operation in operations { + if Operation::get_opid(*operation, ctx) == CuteSchedulerNew1dOp::get_opid_static() { + let scheduler = CuteSchedulerNew1dOp::wrap(*operation); + facts.insert(scheduler.current(ctx), SchedulerFact::Current(*operation)); + facts.insert(scheduler.stride(ctx), SchedulerFact::Stride(*operation)); + } + } + + loop { + let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); + for operation in operations { + if Operation::get_opid(*operation, ctx) != CuteSchedulerAdvanceOp::get_opid_static() { + continue; + } + let advance = CuteSchedulerAdvanceOp::wrap(*operation); + let Some(current_facts) = facts.at(advance.current(ctx)).cloned() else { + continue; + }; + let Some(stride_facts) = facts.at(advance.stride(ctx)).cloned() else { + continue; + }; + for current in ¤t_facts { + let SchedulerFact::Current(root) = current else { + continue; + }; + if stride_facts.contains(&SchedulerFact::Stride(*root)) { + changed |= facts.insert(advance.next(ctx), SchedulerFact::Current(*root)); + } + } + } + if !changed { + break; + } + } + + // A factless scalar is an ordinary, unrelated value. Seeding it only + // after the semantic fixed point lets a valid loop recurrence establish + // itself first, while still poisoning a merge with an arbitrary edge. + seed_unknown_integer_carriers(ctx, operations, &mut facts, 64, SchedulerFact::Unknown); + loop { + let poisoned = + poison_missing_forwarding_sources(ctx, operations, &mut facts, SchedulerFact::Unknown); + let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); + if !poisoned && !propagated { + break; + } + } + facts +} + +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), + )); + } +} + +fn verify_scheduler_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = scheduler_semantic_ids(); + let facts = build_scheduler_facts(ctx, all_ops); + let mut roots = HashSet::new(); + let mut has_work_roots = HashSet::new(); + let mut current_roots = HashSet::new(); + let mut advance_roots = HashSet::new(); + let mut coordinate_roots = HashSet::new(); + let mut has_work_by_root: HashMap, Vec>> = HashMap::new(); + let mut current_by_root: HashMap, Vec>> = HashMap::new(); + let mut advance_by_root: HashMap, Vec>> = HashMap::new(); + let mut coordinates_by_root: HashMap, Vec>> = HashMap::new(); + let mut currents = Vec::new(); + let mut coordinates = Vec::new(); + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if !ids.contains(&opid) { + continue; + } + if operation.deref(ctx).get_parent_block().is_none() { + return Err(invalid(format!( + "scheduler operation `{opid}` is not attached to a block" + ))); + } + let grid = scheduler_grid(ctx, *operation, &opid)?; + if opid == CuteSchedulerNew1dOp::get_opid_static() { + roots.insert(*operation); + } else if opid == CuteSchedulerHasWorkOp::get_opid_static() { + let semantic = CuteSchedulerHasWorkOp::wrap(*operation); + let root = scheduler_root( + &facts, + semantic.current(ctx), + true, + "cute.scheduler_has_work current", + )?; + if CuteSchedulerNew1dOp::wrap(root).tile_grid(ctx) != Some(grid) { + return Err(invalid( + "cute.scheduler_has_work grid differs from its scheduler_new_1d root", + )); + } + if !semantic.has_work(ctx).is_used(ctx) { + return Err(invalid( + "cute.scheduler_has_work result is unused; it must guard the scheduler body", + )); + } + has_work_roots.insert(root); + has_work_by_root.entry(root).or_default().push(*operation); + } else if opid == CuteSchedulerCurrentOp::get_opid_static() { + let semantic = CuteSchedulerCurrentOp::wrap(*operation); + let root = scheduler_root( + &facts, + semantic.current(ctx), + true, + "cute.scheduler_current current", + )?; + if CuteSchedulerNew1dOp::wrap(root).tile_grid(ctx) != Some(grid) { + return Err(invalid( + "cute.scheduler_current grid differs from its scheduler_new_1d root", + )); + } + current_roots.insert(root); + current_by_root.entry(root).or_default().push(*operation); + currents.push(*operation); + } else if opid == CuteWorkTileCoordinatesOp::get_opid_static() { + coordinates.push(*operation); + } else if opid == CuteSchedulerAdvanceOp::get_opid_static() { + let semantic = CuteSchedulerAdvanceOp::wrap(*operation); + let current_root = scheduler_root( + &facts, + semantic.current(ctx), + true, + "cute.scheduler_advance current", + )?; + let stride_root = scheduler_root( + &facts, + semantic.stride(ctx), + false, + "cute.scheduler_advance stride", + )?; + if current_root != stride_root + || CuteSchedulerNew1dOp::wrap(current_root).tile_grid(ctx) != Some(grid) + { + return Err(invalid( + "cute.scheduler_advance must use current, stride, and grid from one scheduler_new_1d", + )); + } + advance_roots.insert(current_root); + advance_by_root + .entry(current_root) + .or_default() + .push(*operation); + } + } + + for current in ¤ts { + let tile = CuteSchedulerCurrentOp::wrap(*current).work_tile(ctx); + let uses = tile.uses(ctx); + if uses.is_empty() + || uses.iter().any(|r#use| { + Operation::get_opid(r#use.user_op(), ctx) + != CuteWorkTileCoordinatesOp::get_opid_static() + || r#use.find_index(ctx) != 0 + }) + { + return Err(invalid( + "cute.scheduler_current work tile must be used only by at least one cute.work_tile_coordinates", + )); + } + } + for coordinates in coordinates { + let tile = CuteWorkTileCoordinatesOp::wrap(coordinates).work_tile(ctx); + let Some(producer) = tile.defining_op() else { + return Err(invalid( + "cute.work_tile_coordinates reached a block argument; a direct scheduler_current result is required", + )); + }; + if Operation::get_opid(producer, ctx) != CuteSchedulerCurrentOp::get_opid_static() + || !currents.contains(&producer) + { + return Err(invalid( + "cute.work_tile_coordinates must use a scheduler_current attached to this module", + )); + } + let root = scheduler_root( + &facts, + CuteSchedulerCurrentOp::wrap(producer).current(ctx), + true, + "cute.work_tile_coordinates scheduler current", + )?; + coordinate_roots.insert(root); + coordinates_by_root + .entry(root) + .or_default() + .push(coordinates); + } + + for root in roots { + if !has_work_roots.contains(&root) + || !current_roots.contains(&root) + || !coordinate_roots.contains(&root) + || !advance_roots.contains(&root) + { + return Err(invalid( + "each cute.scheduler_new_1d needs a complete has_work -> current -> coordinates -> advance story", + )); + } + let has_work = &has_work_by_root[&root]; + let current = ¤t_by_root[&root]; + let coordinates = &coordinates_by_root[&root]; + let advances = &advance_by_root[&root]; + if has_work.len() != 1 + || current.len() != 1 + || coordinates.len() != 1 + || advances.len() != 1 + { + return Err(invalid( + "each scheduler root needs exactly one static has_work/current/coordinates/advance protocol", + )); + } + let Some((guard, work_successor, exit_successor)) = + scheduler_guard(ctx, CuteSchedulerHasWorkOp::wrap(has_work[0]).has_work(ctx)) + else { + return Err(invalid( + "cute.scheduler_has_work must exclusively guard one conditional branch (optionally through mir.not)", + )); + }; + let current_block = current[0] + .deref(ctx) + .get_parent_block() + .expect("attached scheduler current"); + if !operation_reaches(ctx, root, has_work[0]) + || !block_reaches(ctx, work_successor, current_block) + || block_reaches(ctx, exit_successor, current_block) + || !operation_reaches(ctx, current[0], coordinates[0]) + || !operation_reaches(ctx, coordinates[0], advances[0]) + || !operation_reaches(ctx, advances[0], guard) + { + return Err(invalid( + "scheduler CFG must be new -> guarded work current -> coordinates -> advance -> next guard, with the exit edge outside the work body", + )); + } + } + Ok(()) +} + +#[derive(Clone, Copy)] +struct LoadPipelineFacts { + make: Ptr, + stages: u64, + transaction_bytes: u32, +} + +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, + }) +} + +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 + } +} + +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) +} + +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +enum PipelineFact { + Slot(Ptr), + Phase(Ptr), + Unknown(Value), +} + +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" + ))), + } +} + +fn build_pipeline_facts( + ctx: &Context, + operations: &[Ptr], +) -> OrdinaryFacts { + let mut facts = OrdinaryFacts::default(); + for operation in operations { + if Operation::get_opid(*operation, ctx) == CutePipelineStateNewOp::get_opid_static() { + let state = CutePipelineStateNewOp::wrap(*operation); + facts.insert(state.slot(ctx), PipelineFact::Slot(*operation)); + facts.insert(state.phase(ctx), PipelineFact::Phase(*operation)); + } + } + + loop { + let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); + for operation in operations { + if Operation::get_opid(*operation, ctx) != CutePipelineStateAdvanceOp::get_opid_static() + { + continue; + } + let advance = CutePipelineStateAdvanceOp::wrap(*operation); + let Some(slot_facts) = facts.at(advance.slot(ctx)).cloned() else { + continue; + }; + let Some(phase_facts) = facts.at(advance.phase(ctx)).cloned() else { + continue; + }; + for slot in &slot_facts { + let PipelineFact::Slot(root) = slot else { + continue; + }; + if phase_facts.contains(&PipelineFact::Phase(*root)) + && CutePipelineStateNewOp::wrap(*root).state(ctx) == advance.state(ctx) + { + changed |= facts.insert(advance.next_slot(ctx), PipelineFact::Slot(*root)); + changed |= facts.insert(advance.next_phase(ctx), PipelineFact::Phase(*root)); + } + } + } + if !changed { + break; + } + } + + seed_unknown_integer_carriers(ctx, operations, &mut facts, 32, PipelineFact::Unknown); + loop { + let poisoned = + poison_missing_forwarding_sources(ctx, operations, &mut facts, PipelineFact::Unknown); + let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); + if !poisoned && !propagated { + break; + } + } + facts +} + +#[derive(Default)] +struct PipelineStateEvents { + slots: Vec>, + advances: Vec>, + acquires: Vec>, + expects: Vec>, + waits: Vec>, + releases: Vec>, + tails: Vec>, +} + +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", + )) + } +} + +fn verify_pipeline_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = pipeline_semantic_ids(); + let facts = build_pipeline_facts(ctx, all_ops); + let makes: Vec<_> = all_ops + .iter() + .copied() + .filter(|operation| { + Operation::get_opid(*operation, ctx) == CuteTmaLoadPipelineMakeOp::get_opid_static() + }) + .collect(); + let mut init_by_make: HashMap, Vec>> = HashMap::new(); + let mut events: HashMap, PipelineStateEvents> = HashMap::new(); + let mut binding: HashMap, Ptr> = HashMap::new(); + + for make in &makes { + if make.deref(ctx).get_parent_block().is_none() { + return Err(invalid("TMA load-pipeline make is not attached to a block")); + } + let handle = CuteTmaLoadPipelineMakeOp::wrap(*make).pipeline(ctx); + let uses = handle.uses(ctx); + if uses.is_empty() { + return Err(invalid( + "cute.tma_load_pipeline_make has no lifecycle consumer", + )); + } + for r#use in uses { + let user = r#use.user_op(); + let user_id = Operation::get_opid(user, ctx); + if !ids.contains(&user_id) + || r#use.find_index(ctx) != 0 + || pipeline_operand(ctx, user, &user_id) != Some(handle) + { + return Err(invalid(format!( + "TMA load-pipeline handle has an invalid use by `{user_id}`" + ))); + } + } + } + + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if !ids.contains(&opid) { + continue; + } + if operation.deref(ctx).get_parent_block().is_none() { + return Err(invalid(format!( + "pipeline operation `{opid}` is not attached to a block" + ))); + } + if opid == CutePipelineStateNewOp::get_opid_static() { + events.entry(*operation).or_default(); + continue; + } + + let state_root = if opid == CutePipelineStateSlotOp::get_opid_static() { + let semantic = CutePipelineStateSlotOp::wrap(*operation); + let root = pipeline_state_root( + &facts, + semantic.slot(ctx), + true, + "cute.pipeline_state_slot slot", + )?; + if semantic.state(ctx) != CutePipelineStateNewOp::wrap(root).state(ctx) { + return Err(invalid( + "cute.pipeline_state_slot facts differ from its pipeline_state_new root", + )); + } + events.entry(root).or_default().slots.push(*operation); + Some(root) + } else if opid == CutePipelineStateAdvanceOp::get_opid_static() { + let semantic = CutePipelineStateAdvanceOp::wrap(*operation); + let slot_root = pipeline_state_root( + &facts, + semantic.slot(ctx), + true, + "cute.pipeline_state_advance slot", + )?; + let phase_root = pipeline_state_root( + &facts, + semantic.phase(ctx), + false, + "cute.pipeline_state_advance phase", + )?; + if slot_root != phase_root + || semantic.state(ctx) != CutePipelineStateNewOp::wrap(slot_root).state(ctx) + { + return Err(invalid( + "cute.pipeline_state_advance must consume slot and phase from one matching pipeline_state_new", + )); + } + events + .entry(slot_root) + .or_default() + .advances + .push(*operation); + Some(slot_root) + } else if opid == CuteTmaLoadPipelineMakeOp::get_opid_static() { + None + } else if opid == CuteTmaLoadPipelineInitOp::get_opid_static() { + let pipeline = resolve_load_pipeline( + ctx, + CuteTmaLoadPipelineInitOp::wrap(*operation).pipeline(ctx), + )?; + if !makes.contains(&pipeline.make) { + return Err(invalid( + "pipeline init refers to a make outside its semantic scope", + )); + } + init_by_make + .entry(pipeline.make) + .or_default() + .push(*operation); + None + } else { + let state = pipeline_state_for_op(ctx, *operation, &opid) + .ok_or_else(|| invalid(format!("pipeline operation `{opid}` has no state")))?; + let operation_ref = operation.deref(ctx); + let slot_root = pipeline_state_root( + &facts, + operation_ref.get_operand(1), + true, + &format!("`{opid}` slot"), + )?; + let needs_phase = opid == CutePipelineProducerAcquireOp::get_opid_static() + || opid == CutePipelineConsumerWaitOp::get_opid_static() + || opid == CutePipelineProducerTailOp::get_opid_static(); + if needs_phase { + let phase_root = pipeline_state_root( + &facts, + operation_ref.get_operand(2), + false, + &format!("`{opid}` phase"), + )?; + if phase_root != slot_root { + return Err(invalid(format!( + "`{opid}` mixes slot and phase from different pipeline states" + ))); + } + } + if CutePipelineStateNewOp::wrap(slot_root).state(ctx) != Some(state) { + return Err(invalid(format!( + "`{opid}` state facts differ from its pipeline_state_new root" + ))); + } + + let handle = pipeline_operand(ctx, *operation, &opid) + .ok_or_else(|| invalid(format!("`{opid}` has no load-pipeline handle")))?; + let pipeline = resolve_load_pipeline(ctx, handle)?; + if !makes.contains(&pipeline.make) { + return Err(invalid( + "pipeline lifecycle operation refers to a make outside its semantic scope", + )); + } + if state.stages != pipeline.stages { + return Err(invalid(format!( + "pipeline state on `{opid}` has {} stages but its handle has {}", + state.stages, pipeline.stages + ))); + } + if let Some(previous) = binding.insert(slot_root, pipeline.make) + && previous != pipeline.make + { + return Err(invalid( + "one ordinary pipeline state is used with multiple load-pipeline handles", + )); + } + + let root_events = events.entry(slot_root).or_default(); + if opid == CutePipelineProducerAcquireOp::get_opid_static() { + root_events.acquires.push(*operation); + } else if opid == CutePipelineProducerExpectTxOp::get_opid_static() { + root_events.expects.push(*operation); + let copied = expected_tma_transaction_bytes(ctx, *operation)?; + if copied != u64::from(pipeline.transaction_bytes) { + return Err(invalid(format!( + "cute.pipeline_producer_expect_tx promises {} bytes, but its direct TMA copies move {copied} bytes", + pipeline.transaction_bytes + ))); + } + } else if opid == CutePipelineConsumerWaitOp::get_opid_static() { + root_events.waits.push(*operation); + } else if opid == CutePipelineConsumerReleaseOp::get_opid_static() { + root_events.releases.push(*operation); + } else if opid == CutePipelineProducerTailOp::get_opid_static() { + root_events.tails.push(*operation); + } + Some(slot_root) + }; + + if let Some(root) = state_root { + let root_state = CutePipelineStateNewOp::wrap(root) + .state(ctx) + .expect("locally verified pipeline state root"); + if let Some(state) = pipeline_state_for_op(ctx, *operation, &opid) + && state != root_state + { + return Err(invalid(format!( + "`{opid}` does not carry its pipeline-state root's role and stage count" + ))); + } + } + } + + let mut roles_by_make: HashMap, (usize, usize)> = HashMap::new(); + for (root, root_events) in &events { + let state = CutePipelineStateNewOp::wrap(*root) + .state(ctx) + .expect("locally verified pipeline state root"); + let Some(make) = binding.get(root).copied() else { + return Err(invalid( + "cute.pipeline_state_new is not connected to any load-pipeline lifecycle", + )); + }; + if root_events.slots.is_empty() || root_events.advances.is_empty() { + return Err(invalid( + "each pipeline state needs state_slot and state_advance operations", + )); + } + let roles = roles_by_make.entry(make).or_default(); + match state.role { + CutePipelineRoleAttr::Producer => { + roles.0 += 1; + if root_events.acquires.len() != 1 + || root_events.expects.len() != 1 + || root_events.advances.len() != 1 + || root_events.slots.len() != 1 + || root_events.tails.len() != 1 + || !root_events.waits.is_empty() + || !root_events.releases.is_empty() + { + return Err(invalid( + "producer pipeline state needs acquire, expect_tx, advance, and exactly one tail, with no consumer events", + )); + } + let acquire = root_events.acquires[0]; + let expect = root_events.expects[0]; + let advance = root_events.advances[0]; + let tail = root_events.tails[0]; + require_producer_pipeline_order(ctx, *root, acquire, expect, advance, tail)?; + let barrier = CutePipelineProducerExpectTxOp::wrap(expect).completion_barrier(ctx); + for r#use in barrier.uses(ctx) { + let copy = r#use.user_op(); + if !operation_dominates(ctx, expect, copy) + || !operation_dominates(ctx, copy, advance) + { + return Err(invalid( + "every TMA copy must be CFG-ordered after expect_tx and before producer state advance", + )); + } + } + } + CutePipelineRoleAttr::Consumer => { + roles.1 += 1; + if root_events.waits.len() != 1 + || root_events.releases.len() != 1 + || root_events.advances.len() != 1 + || root_events.slots.len() != 1 + || !root_events.acquires.is_empty() + || !root_events.expects.is_empty() + || !root_events.tails.is_empty() + { + return Err(invalid( + "consumer pipeline state needs wait, release, and advance, with no producer events", + )); + } + let wait = root_events.waits[0]; + let release = root_events.releases[0]; + let advance = root_events.advances[0]; + if !operation_dominates(ctx, *root, wait) + || !operation_dominates(ctx, wait, release) + || !operation_dominates(ctx, release, advance) + || !operation_reaches(ctx, advance, wait) + { + return Err(invalid( + "consumer pipeline CFG must order state_new -> wait -> release -> advance -> next wait", + )); + } + } + } + } + for make in makes { + if init_by_make + .get(&make) + .is_none_or(|initializers| initializers.len() != 1) + { + return Err(invalid( + "each tma_load_pipeline_make needs exactly one initialization", + )); + } + if roles_by_make.get(&make).copied() != Some((1, 1)) { + return Err(invalid( + "each tma_load_pipeline_make needs exactly one complete producer state and one complete consumer state", + )); + } + let initializer = init_by_make[&make][0]; + for (root, bound_make) in &binding { + if *bound_make == make && !operation_dominates(ctx, initializer, *root) { + return Err(invalid( + "load-pipeline initialization must dominate both producer and consumer state creation", + )); + } + } + } + Ok(()) +} + +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) +} + +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(()) +} + +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(()) +} + +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 + } +} + +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +enum MmaFact { + Lane(Ptr), + SmemBase(Ptr), + SmemCapacity(Ptr), + StageScales(Ptr), + SlicedScales(Ptr), + AFragment(Ptr), + BBase(Ptr), + BCapacity(Ptr), + BWarpN(Ptr), + BKHalf(Ptr), + Accumulator(Ptr), + Scalar(Value), + Unknown(Value), +} + +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")) +} + +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" + ))), + } +} + +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) +} + +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" + ))), + } +} + +fn build_mma_facts(ctx: &Context, operations: &[Ptr]) -> OrdinaryFacts { + let mut facts = OrdinaryFacts::default(); + for operation in operations { + let opid = Operation::get_opid(*operation, ctx); + let operation_ref = operation.deref(ctx); + if opid == CuteSmemTensorOverlayOp::get_opid_static() { + let overlay = CuteSmemTensorOverlayOp::wrap(*operation); + facts.insert(overlay.base(ctx), MmaFact::SmemBase(*operation)); + facts.insert(overlay.capacity(ctx), MmaFact::SmemCapacity(*operation)); + } else if opid == CuteTiledMmaSliceOp::get_opid_static() { + facts.insert( + CuteTiledMmaSliceOp::wrap(*operation).sliced_lane(ctx), + MmaFact::Lane(*operation), + ); + } else if opid == CuteFragmentFillOp::get_opid_static() { + facts.insert( + CuteFragmentFillOp::wrap(*operation).fragment(ctx), + MmaFact::Accumulator(*operation), + ); + } else if opid == CuteMmaLoadScalesOp::get_opid_static() { + facts.insert( + CuteMmaLoadScalesOp::wrap(*operation).scales(ctx), + MmaFact::StageScales(*operation), + ); + } else if opid == CuteFragmentSliceKOp::get_opid_static() { + facts.insert( + CuteFragmentSliceKOp::wrap(*operation).scales(ctx), + MmaFact::SlicedScales(*operation), + ); + } else if opid == CuteMmaLoadAOp::get_opid_static() { + facts.insert( + CuteMmaLoadAOp::wrap(*operation).fragment(ctx), + MmaFact::AFragment(*operation), + ); + } else if opid == CuteMmaPartitionBOp::get_opid_static() { + facts.insert(operation_ref.get_result(0), MmaFact::BBase(*operation)); + facts.insert(operation_ref.get_result(1), MmaFact::BCapacity(*operation)); + facts.insert(operation_ref.get_result(2), MmaFact::BWarpN(*operation)); + facts.insert(operation_ref.get_result(3), MmaFact::BKHalf(*operation)); + } + } + + loop { + let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); + for operation in operations { + if Operation::get_opid(*operation, ctx) != CuteTiledGemmOp::get_opid_static() { + continue; + } + let operation_ref = operation.deref(ctx); + let Some(accumulator_facts) = facts.at(operation_ref.get_operand(7)).cloned() else { + continue; + }; + for accumulator in accumulator_facts { + if let MmaFact::Accumulator(root) = accumulator { + changed |= facts.insert( + CuteTiledGemmOp::wrap(*operation).accumulator(ctx), + MmaFact::Accumulator(root), + ); + } + } + } + if !changed { + break; + } + } + + seed_unknown_integer_carriers(ctx, operations, &mut facts, 32, MmaFact::Scalar); + seed_unknown_integer_carriers(ctx, operations, &mut facts, 64, MmaFact::Scalar); + loop { + let poisoned = + poison_missing_forwarding_sources(ctx, operations, &mut facts, MmaFact::Unknown); + let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); + if !poisoned && !propagated { + break; + } + } + facts +} + +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)) +} + +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, + } +} + +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, + ) + }) + }) + }) +} + +fn mma_ordinary_use_forwards_fact( + ctx: &Context, + facts: &OrdinaryFacts, + key: &CarrierKey, + fact: MmaFact, + user: Ptr, + operand_index: usize, +) -> bool { + let opid = Operation::get_opid(user, ctx); + let operation = user.deref(ctx); + if opid == MirCastOp::get_opid_static() + && operand_index == 0 + && is_semantic_pointer_carrier_cast(ctx, user) + && operation.get_num_results() == 1 + { + return mma_fact_at_path(facts, operation.get_result(0), &key.path, fact); + } + if (opid == MirConstructStructOp::get_opid_static() + || opid == MirConstructTupleOp::get_opid_static()) + && operation.get_num_results() == 1 + { + let mut destination_path = vec![operand_index as u32]; + destination_path.extend_from_slice(&key.path); + return mma_fact_at_path(facts, operation.get_result(0), &destination_path, fact); + } + if opid == MirInsertFieldOp::get_opid_static() + && operation.get_num_operands() == 2 + && operation.get_num_results() == 1 + { + let Some(index) = MirInsertFieldOp::new(user).get_attr_insert_index(ctx) else { + return false; + }; + if operand_index == 0 { + if key.path.first().copied() == Some(index.0) { + return true; + } + return mma_fact_at_path(facts, operation.get_result(0), &key.path, fact); + } + if operand_index == 1 { + let mut destination_path = vec![index.0]; + destination_path.extend_from_slice(&key.path); + return mma_fact_at_path(facts, operation.get_result(0), &destination_path, fact); + } + return false; + } + if opid == MirExtractFieldOp::get_opid_static() + && operand_index == 0 + && operation.get_num_results() == 1 + { + let Some(index) = MirExtractFieldOp::new(user).get_attr_index(ctx) else { + return false; + }; + if key.path.first().copied() != Some(index.0) { + return true; + } + return mma_fact_at_path(facts, operation.get_result(0), &key.path[1..], fact); + } + if opid == MirStoreOp::get_opid_static() && operation.get_num_operands() == 2 { + let store = MirStoreOp::new(user); + if operand_index == 0 { + return local_cell(ctx, store.address_opd(ctx)) + .is_some_and(|(root, _)| local_cell_is_closed(ctx, root)); + } + if operand_index == 1 + && let Some((root, mut destination_path)) = local_cell(ctx, store.address_opd(ctx)) + && local_cell_is_closed(ctx, root) + { + destination_path.extend_from_slice(&key.path); + return mma_fact_at_path(facts, root, &destination_path, fact); + } + return false; + } + if opid == MirLoadOp::get_opid_static() + && operand_index == 0 + && operation.get_num_results() == 1 + { + let load = MirLoadOp::new(user); + let Some((root, cell_path)) = local_cell(ctx, load.address_opd(ctx)) else { + return false; + }; + if !local_cell_is_closed(ctx, root) { + return false; + } + let Some(result_path) = key.path.strip_prefix(cell_path.as_slice()) else { + return true; + }; + return mma_fact_at_path(facts, operation.get_result(0), result_path, fact); + } + if opid == MirFieldAddrOp::get_opid_static() + && operand_index == 0 + && operation.get_num_results() == 1 + && MirFieldAddrOp::new(user) + .get_attr_field_index(ctx) + .is_some() + { + return local_cell(ctx, operation.get_result(0)) + .is_some_and(|(root, _)| local_cell_is_closed(ctx, root)); + } + if opid == MirCondBranchOp::get_opid_static() && operand_index == 0 { + return false; + } + mma_branch_forwards_fact(ctx, facts, user, key.value, &key.path, fact) +} + +fn verify_mma_use_closure( + ctx: &Context, + operations: &[Ptr], + facts: &OrdinaryFacts, +) -> Result { + for (key, found) in &facts.facts { + for fact in found.iter().copied() { + if matches!(fact, MmaFact::Scalar(_) | MmaFact::Unknown(_)) { + continue; + } + for r#use in key.value.uses(ctx) { + let user = r#use.user_op(); + let operand_index = r#use.find_index(ctx); + if operations.contains(&user) + && (mma_semantic_consumer(ctx, user, operand_index, fact) + || mma_ordinary_use_forwards_fact( + ctx, + facts, + key, + fact, + user, + operand_index, + )) + { + continue; + } + return Err(invalid(format!( + "shared-MMA carrier has an unsupported outgoing use by `{}` operand {operand_index}", + Operation::get_opid(user, ctx) + ))); + } + } + } + Ok(()) +} + +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(()) +} + +fn verify_smem_mma_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = smem_mma_semantic_ids(); + let facts = build_mma_facts(ctx, all_ops); + let mut overlays = HashSet::new(); + let mut used_overlays = HashSet::new(); + let mut slices = HashSet::new(); + let mut used_slices = HashSet::new(); + let mut fills = HashSet::new(); + let mut used_fills = HashSet::new(); + let mut scale_loads = HashSet::new(); + let mut used_scale_loads = HashSet::new(); + let mut scale_slices = HashSet::new(); + let mut used_scale_slices = HashSet::new(); + let mut a_loads = HashSet::new(); + let mut used_a_loads = HashSet::new(); + let mut b_partitions = HashSet::new(); + let mut used_b_partitions = HashSet::new(); + let mut gemms = 0usize; + + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if !ids.contains(&opid) { + continue; + } + if operation.deref(ctx).get_parent_block().is_none() { + return Err(invalid(format!( + "shared/MMA operation `{opid}` is not attached to a block" + ))); + } + let operation_ref = operation.deref(ctx); + if opid == CuteSmemTensorOverlayOp::get_opid_static() { + overlays.insert(*operation); + } else if opid == CuteTiledMmaSliceOp::get_opid_static() { + slices.insert(*operation); + } else if opid == CuteFragmentFillOp::get_opid_static() { + fills.insert(*operation); + } else if opid == CuteMmaLoadScalesOp::get_opid_static() { + scale_loads.insert(*operation); + let semantic = CuteMmaLoadScalesOp::wrap(*operation); + let lane = require_mma_lane( + &facts, + operation_ref.get_operand(0), + "cute.mma_load_scales lane", + )?; + require_same_mma_plan(ctx, lane, *operation, "mma scale load")?; + used_slices.insert(lane); + let scale_a = require_smem_overlay_pair( + &facts, + operation_ref.get_operand(1), + operation_ref.get_operand(2), + "cute.mma_load_scales A", + )?; + let scale_b = require_smem_overlay_pair( + &facts, + operation_ref.get_operand(3), + operation_ref.get_operand(4), + "cute.mma_load_scales B", + )?; + if semantic.scale_a_view(ctx) != CuteSmemTensorOverlayOp::wrap(scale_a).view(ctx) + || semantic.scale_b_view(ctx) != CuteSmemTensorOverlayOp::wrap(scale_b).view(ctx) + { + return Err(invalid( + "cute.mma_load_scales typed views do not match its exact shared overlays", + )); + } + used_overlays.extend([scale_a, scale_b]); + } else if opid == CuteFragmentSliceKOp::get_opid_static() { + scale_slices.insert(*operation); + let load = match one_mma_fact( + &facts, + operation_ref.get_operand(0), + "cute.fragment_slice_k stage scales", + )? { + MmaFact::StageScales(load) => load, + _ => { + return Err(invalid( + "cute.fragment_slice_k must consume cute.mma_load_scales", + )); + } + }; + require_same_mma_plan(ctx, load, *operation, "scale K slice")?; + used_scale_loads.insert(load); + } else if opid == CuteMmaLoadAOp::get_opid_static() { + a_loads.insert(*operation); + let semantic = CuteMmaLoadAOp::wrap(*operation); + let lane = + require_mma_lane(&facts, operation_ref.get_operand(0), "cute.mma_load_a lane")?; + require_same_mma_plan(ctx, lane, *operation, "MMA A load")?; + used_slices.insert(lane); + let overlay = require_smem_overlay_pair( + &facts, + operation_ref.get_operand(1), + operation_ref.get_operand(2), + "cute.mma_load_a", + )?; + if semantic.view(ctx) != CuteSmemTensorOverlayOp::wrap(overlay).view(ctx) { + return Err(invalid( + "cute.mma_load_a typed view does not match its exact shared overlay", + )); + } + used_overlays.insert(overlay); + } else if opid == CuteMmaPartitionBOp::get_opid_static() { + b_partitions.insert(*operation); + let semantic = CuteMmaPartitionBOp::wrap(*operation); + let overlay = require_smem_overlay_pair( + &facts, + operation_ref.get_operand(0), + operation_ref.get_operand(1), + "cute.mma_partition_b", + )?; + if semantic.view(ctx) != CuteSmemTensorOverlayOp::wrap(overlay).view(ctx) { + return Err(invalid( + "cute.mma_partition_b typed view does not match its exact shared overlay", + )); + } + used_overlays.insert(overlay); + } else if opid == CuteTiledGemmOp::get_opid_static() { + gemms += 1; + let semantic = CuteTiledGemmOp::wrap(*operation); + let lane = + require_mma_lane(&facts, operation_ref.get_operand(0), "cute.tiled_gemm lane")?; + let a = match one_mma_fact( + &facts, + operation_ref.get_operand(1), + "cute.tiled_gemm A fragment", + )? { + MmaFact::AFragment(load) => load, + _ => return Err(invalid("cute.tiled_gemm A must come from cute.mma_load_a")), + }; + let b = match ( + one_mma_fact(&facts, operation_ref.get_operand(2), "tiled_gemm B base")?, + one_mma_fact( + &facts, + operation_ref.get_operand(3), + "tiled_gemm B capacity", + )?, + one_mma_fact(&facts, operation_ref.get_operand(4), "tiled_gemm B warp-N")?, + one_mma_fact(&facts, operation_ref.get_operand(5), "tiled_gemm B K-half")?, + ) { + ( + MmaFact::BBase(base), + MmaFact::BCapacity(capacity), + MmaFact::BWarpN(warp_n), + MmaFact::BKHalf(k_half), + ) if base == capacity && base == warp_n && base == k_half => base, + _ => { + return Err(invalid( + "cute.tiled_gemm B carriers must be the four matching results of one mma_partition_b", + )); + } + }; + let scales = match one_mma_fact( + &facts, + operation_ref.get_operand(6), + "cute.tiled_gemm scales", + )? { + MmaFact::SlicedScales(slice) => slice, + _ => { + return Err(invalid( + "cute.tiled_gemm scales must come from cute.fragment_slice_k", + )); + } + }; + let fill = match one_mma_fact( + &facts, + operation_ref.get_operand(7), + "cute.tiled_gemm accumulator", + )? { + MmaFact::Accumulator(fill) => fill, + _ => { + return Err(invalid( + "cute.tiled_gemm accumulator must be the closed recurrence seeded by fragment_fill", + )); + } + }; + for producer in [lane, a, b, scales, fill] { + require_same_mma_plan(ctx, producer, *operation, "tiled GEMM spine")?; + } + let scale_load = match one_mma_fact( + &facts, + scales.deref(ctx).get_operand(0), + "tiled GEMM scale-load origin", + )? { + MmaFact::StageScales(load) => load, + _ => unreachable!("fragment_slice_k was verified above"), + }; + let scale_load_ref = scale_load.deref(ctx); + let a_ref = a.deref(ctx); + let b_ref = b.deref(ctx); + let scale_lane = require_mma_lane( + &facts, + scale_load_ref.get_operand(0), + "mma_load_scales lane", + )?; + let a_lane = require_mma_lane(&facts, a_ref.get_operand(0), "mma_load_a lane")?; + if scale_lane != lane || a_lane != lane { + return Err(invalid( + "tiled GEMM, scale load, and A load must use one tiled-MMA lane slice", + )); + } + require_same_mma_scalar( + &facts, + scale_load_ref.get_operand(5), + a_ref.get_operand(3), + "shared-MMA warp-M", + )?; + require_same_mma_scalar( + &facts, + scale_load_ref.get_operand(6), + b_ref.get_operand(2), + "shared-MMA warp-N", + )?; + require_same_mma_scalar( + &facts, + scales.deref(ctx).get_operand(1), + a_ref.get_operand(4), + "shared-MMA K-half", + )?; + require_same_mma_scalar( + &facts, + a_ref.get_operand(4), + b_ref.get_operand(3), + "shared-MMA K-half", + )?; + if semantic.b_view(ctx) != CuteMmaPartitionBOp::wrap(b).view(ctx) { + return Err(invalid( + "cute.tiled_gemm B view does not match its mma_partition_b producer", + )); + } + used_slices.insert(lane); + used_a_loads.insert(a); + used_b_partitions.insert(b); + used_scale_slices.insert(scales); + used_fills.insert(fill); + } + } + + let found = overlays.len() + + slices.len() + + fills.len() + + scale_loads.len() + + scale_slices.len() + + a_loads.len() + + b_partitions.len() + + gemms; + if found == 0 { + return Ok(()); + } + verify_mma_use_closure(ctx, all_ops, &facts)?; + if gemms == 0 + || overlays != used_overlays + || slices != used_slices + || fills != used_fills + || scale_loads != used_scale_loads + || scale_slices != used_scale_slices + || a_loads != used_a_loads + || b_partitions != used_b_partitions + { + return Err(invalid( + "shared-MMA story is incomplete or has an unconsumed overlay/slice/load/partition/fill producer", + )); + } + Ok(()) +} + +/// 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. +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; + } +} + +fn is_epilogue_pointer_identity_cast(ctx: &Context, operation: Ptr) -> bool { + is_semantic_pointer_carrier_cast(ctx, operation) +} + +fn is_epilogue_writer_plumbing(ctx: &Context, operation: Ptr) -> bool { + let opid = Operation::get_opid(operation, ctx); + opid == MirGotoOp::get_opid_static() + || opid == MirStorageLiveOp::get_opid_static() + || opid == MirStorageDeadOp::get_opid_static() + || opid == MirExtractFieldOp::get_opid_static() + || is_epilogue_pointer_identity_cast(ctx, operation) + || is_closed_local_cell_load(ctx, operation) +} + +/// A non-volatile `mir.load` from a closed compiler-created local cell only +/// re-materializes carrier state (e.g. the epilogue handle struct spilled to +/// an alloca slot); it cannot touch tile memory, so it is admissible plumbing +/// between the writer-protocol anchor operations. +fn is_closed_local_cell_load(ctx: &Context, operation: Ptr) -> bool { + if Operation::get_opid(operation, ctx) != MirLoadOp::get_opid_static() { + return false; + } + let load = MirLoadOp::new(operation); + { + let operation_ref = operation.deref(ctx); + if operation_ref.get_num_operands() != 1 || operation_ref.get_num_results() != 1 { + return false; + } + } + if load.is_volatile(ctx) { + return false; + } + let Some((root, _path)) = local_cell(ctx, load.address_opd(ctx)) else { + return false; + }; + local_cell_is_closed(ctx, root) +} + +/// Accept only the dead immutable reference Rust leaves around a no-inline +/// call on the fieldless, zero-sized TMA store-pipeline token. +fn is_dead_epilogue_store_pipeline_ref(ctx: &Context, operation: Ptr) -> bool { + if Operation::get_opid(operation, ctx) != MirRefOp::get_opid_static() { + return false; + } + let reference = MirRefOp::new(operation); + if reference.verify(ctx).is_err() + || reference + .get_attr_mutable(ctx) + .is_none_or(|mutable| mutable.0) + { + return false; + } + let operation_ref = operation.deref(ctx); + let referent = operation_ref.get_operand(0); + let result = operation_ref.get_result(0); + if result.is_used(ctx) { + return false; + } + + let referent_ty = referent.get_type(ctx); + let referent_ty_ref = referent_ty.deref(ctx); + let Some(store_pipeline) = referent_ty_ref.downcast_ref::() else { + return false; + }; + if store_pipeline.name != "TmaStorePipeline" + || !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 +} + +fn is_epilogue_issuer_plumbing(ctx: &Context, operation: Ptr) -> bool { + let opid = Operation::get_opid(operation, ctx); + is_epilogue_writer_plumbing(ctx, operation) + || opid == MirUndefOp::get_opid_static() + || opid == MirInsertFieldOp::get_opid_static() + || opid == MirConstructStructOp::get_opid_static() + || opid == MirConstructTupleOp::get_opid_static() + || opid == MirConstantOp::get_opid_static() + || opid == MirAddOp::get_opid_static() + || is_dead_epilogue_store_pipeline_ref(ctx, operation) +} + +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", + )) +} + +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(()) +} + +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()) +} + +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(()) +} + +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", + )) + } +} + +#[derive(Clone, Copy)] +struct EpilogueStoreFacts { + operation: Ptr, + source_view: Ptr, + destination_view: Ptr, +} + +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, + }) +} + +fn verify_epilogue_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = epilogue_semantic_ids(); + let mut overlays = Vec::new(); + let mut slices = Vec::new(); + let mut fragments = Vec::new(); + let mut syncs = Vec::new(); + let mut halves = Vec::new(); + let mut acquires = Vec::new(); + let mut commits = Vec::new(); + let mut tails = Vec::new(); + let mut stores = Vec::new(); + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if opid == ids[0] { + overlays.push(*operation); + } else if opid == ids[1] { + slices.push(*operation); + } else if opid == ids[2] { + fragments.push(*operation); + } else if opid == ids[3] { + syncs.push(*operation); + } else if opid == ids[4] { + halves.push(*operation); + } else if opid == ids[5] { + acquires.push(*operation); + } else if opid == ids[6] { + commits.push(*operation); + } else if opid == ids[7] { + tails.push(*operation); + } else if opid == ids[8] { + stores.push(*operation); + } + } + + let found = overlays.len() + + slices.len() + + fragments.len() + + syncs.len() + + halves.len() + + acquires.len() + + commits.len() + + tails.len() + + stores.len(); + if found == 0 { + return Ok(()); + } + if overlays.len() != 1 + || slices.len() != 1 + || fragments.len() != 1 + || syncs.len() != 2 + || halves.len() != 2 + || acquires.len() != 1 + || commits.len() != 1 + || tails.len() != 1 + || stores.len() != 2 + { + return Err(invalid(format!( + "epilogue story needs exactly 1 overlay, 1 warp slice, 1 fragment store, 2 syncs, 2 halves, 1 acquire, 2 TMA stores, 1 commit, and 1 tail; found {}, {}, {}, {}, {}, {}, {}, {}, {}", + overlays.len(), + slices.len(), + fragments.len(), + syncs.len(), + halves.len(), + acquires.len(), + stores.len(), + commits.len(), + tails.len() + ))); + } + + let overlay = overlays[0]; + let warp_slice = slices[0]; + let store_fragment = fragments[0]; + let acquire = acquires[0]; + let commit = commits[0]; + let tail = tails[0]; + let overlay_view = CuteEpilogueSmemOverlayOp::wrap(overlay); + let warp_view = CuteEpilogueWarpSliceOp::wrap(warp_slice); + let fragment = CuteEpilogueStoreFragmentOp::wrap(store_fragment); + let overlay_tile = overlay_view + .tile(ctx) + .expect("locally verified epilogue overlay has a tile"); + if warp_view.tile(ctx).as_ref() != Some(&overlay_tile) + || fragment.tile(ctx).as_ref() != Some(&overlay_tile) + { + return Err(invalid( + "epilogue overlay, warp slice, and fragment store must carry the same tile plan", + )); + } + let mma_facts = build_mma_facts(ctx, all_ops); + let accumulator_seed = match one_mma_fact( + &mma_facts, + fragment.accumulator(ctx), + "epilogue fragment accumulator", + )? { + MmaFact::Accumulator(seed) => seed, + _ => { + return Err(invalid( + "epilogue fragment accumulator must be the closed tiled-GEMM recurrence", + )); + } + }; + if CuteFragmentFillOp::wrap(accumulator_seed).plan(ctx) + != Some(overlay_tile.plan.tiled_mma.clone()) + { + return Err(invalid( + "epilogue accumulator plan differs from its shared-MMA story", + )); + } + require_same_epilogue_carrier( + ctx, + warp_view.input_base(ctx), + overlay_view.base(ctx), + "epilogue warp slice must use the base produced by the epilogue overlay", + )?; + require_same_epilogue_carrier( + ctx, + fragment.base(ctx), + warp_view.base(ctx), + "epilogue fragment store must use the base produced by the warp slice", + )?; + require_same_epilogue_carrier( + ctx, + fragment.warp_id(ctx), + warp_view.warp_id(ctx), + "epilogue fragment store must use the warp ID produced by the warp slice", + )?; + require_same_epilogue_carrier( + ctx, + fragment.lane(ctx), + warp_view.lane(ctx), + "epilogue fragment store must use the lane produced by the warp slice", + )?; + + let reusable_sync = syncs + .iter() + .copied() + .find(|operation| { + CuteEpilogueSyncOp::wrap(*operation).phase(ctx) + == Some(CuteEpilogueSyncPhaseAttr::Reusable) + }) + .ok_or_else(|| invalid("epilogue story is missing the Reusable sync"))?; + let ready_sync = syncs + .iter() + .copied() + .find(|operation| { + CuteEpilogueSyncOp::wrap(*operation).phase(ctx) + == Some(CuteEpilogueSyncPhaseAttr::ReadyForTma) + }) + .ok_or_else(|| invalid("epilogue story is missing the ReadyForTma sync"))?; + if reusable_sync == ready_sync { + return Err(invalid( + "epilogue story needs one distinct sync for each phase", + )); + } + for sync in [reusable_sync, ready_sync] { + let sync = CuteEpilogueSyncOp::wrap(sync); + if sync.tile(ctx).as_ref() != Some(&overlay_tile) { + return Err(invalid( + "both epilogue syncs must carry the overlay tile plan", + )); + } + require_same_epilogue_carrier( + ctx, + sync.base(ctx), + overlay_view.base(ctx), + "both epilogue syncs must use the base produced by the epilogue overlay", + )?; + } + + let writer_range = require_linear_protocol_order( + ctx, + "epilogue writer protocol", + &[reusable_sync, store_fragment, ready_sync], + )?; + for operation in &writer_range { + let expected = [reusable_sync, store_fragment, ready_sync].contains(operation); + if !expected && !is_epilogue_writer_plumbing(ctx, *operation) { + return Err(invalid(format!( + "epilogue writer protocol has interposed non-carrier operation `{}`", + Operation::get_opid(*operation, ctx) + ))); + } + } + let writer_events: Vec<_> = writer_range + .iter() + .copied() + .filter(|operation| { + let opid = Operation::get_opid(*operation, ctx); + opid.dialect.to_string() == crate::CUTE_DIALECT_NAME + || opid == MirCallOp::get_opid_static() + }) + .collect(); + if writer_events != [reusable_sync, store_fragment, ready_sync] { + return Err(invalid( + "epilogue writer protocol must be exactly Reusable -> fragment store -> ReadyForTma, whose semantic boundary includes publication to the TMA proxy", + )); + } + + let half_zero = halves + .iter() + .copied() + .find(|operation| { + CuteEpilogueHalfOp::wrap(*operation) + .half(ctx) + .is_some_and(|half| half.0 == 0) + }) + .ok_or_else(|| invalid("epilogue story is missing half 0"))?; + let half_one = halves + .iter() + .copied() + .find(|operation| { + CuteEpilogueHalfOp::wrap(*operation) + .half(ctx) + .is_some_and(|half| half.0 == 1) + }) + .ok_or_else(|| invalid("epilogue story is missing half 1"))?; + if half_zero == half_one { + return Err(invalid("epilogue story needs two distinct shared halves")); + } + for half in [half_zero, half_one] { + let half = CuteEpilogueHalfOp::wrap(half); + if half.tile(ctx).as_ref() != Some(&overlay_tile) { + return Err(invalid( + "both epilogue halves must carry the overlay tile plan", + )); + } + require_same_epilogue_carrier( + ctx, + half.full_base(ctx), + overlay_view.base(ctx), + "both epilogue halves must use the base produced by the epilogue overlay", + )?; + } + + let store_zero = epilogue_store_facts(ctx, stores[0], all_ops)?; + let store_one = epilogue_store_facts(ctx, stores[1], all_ops)?; + for (half, store) in [(half_zero, store_zero), (half_one, store_one)] { + let half = CuteEpilogueHalfOp::wrap(half); + let source = CuteTmaSmemViewOp::wrap(store.source_view); + require_same_epilogue_carrier( + ctx, + source.base(ctx), + half.half_base(ctx), + "each semantic TMA store source base must come from the matching epilogue half through only the compiler's SSA carrier adapter", + )?; + require_same_epilogue_carrier( + ctx, + source.capacity(ctx), + half.capacity(ctx), + "each semantic TMA store source capacity must come from the matching epilogue half through only the compiler's SSA carrier adapter", + )?; + } + if store_zero.operation.deref(ctx).get_operand(1).get_type(ctx) + != store_one.operation.deref(ctx).get_operand(1).get_type(ctx) + { + return Err(invalid( + "both epilogue stores must target the same typed output tile", + )); + } + let destination_zero = CuteTmaGmemViewOp::wrap(store_zero.destination_view); + let destination_one = CuteTmaGmemViewOp::wrap(store_one.destination_view); + require_same_epilogue_carrier( + ctx, + destination_one.descriptor(ctx), + destination_zero.descriptor(ctx), + "both epilogue halves must target the same TMA destination descriptor", + )?; + let semantic_zero = CuteTmaStore2dSemanticOp::wrap(store_zero.operation); + let semantic_one = CuteTmaStore2dSemanticOp::wrap(store_one.operation); + require_same_epilogue_carrier( + ctx, + semantic_one.tile_row(ctx), + semantic_zero.tile_row(ctx), + "both epilogue halves must target the same output tile row", + )?; + require_next_epilogue_column( + ctx, + semantic_zero.tile_column(ctx), + semantic_one.tile_column(ctx), + )?; + + let issuer_range = require_linear_protocol_order( + ctx, + "epilogue TMA issuer protocol", + &[ + half_zero, + store_zero.operation, + half_one, + store_one.operation, + commit, + ], + )?; + for operation in issuer_range { + let opid = Operation::get_opid(operation, ctx); + let allowed_tma_view = opid == CuteTmaGmemViewOp::get_opid_static() + || opid == CuteTmaSmemViewOp::get_opid_static(); + let expected = [ + half_zero, + store_zero.operation, + half_one, + store_one.operation, + commit, + ] + .contains(&operation); + if !allowed_tma_view && !expected && !is_epilogue_issuer_plumbing(ctx, operation) { + return Err(invalid(format!( + "epilogue TMA issuer protocol has interposed operation `{opid}`" + ))); + } + } + + let acquire_pipeline = CuteTmaStoreAcquireOp::wrap(acquire).pipeline(ctx); + let commit_pipeline = CuteTmaStoreCommitOp::wrap(commit).pipeline(ctx); + let tail_pipeline = CuteTmaStoreTailOp::wrap(tail).pipeline(ctx); + if acquire_pipeline != commit_pipeline || acquire_pipeline != tail_pipeline { + return Err(invalid( + "epilogue acquire, commit, and tail must carry the same TMA store-pipeline facts", + )); + } + + require_epilogue_cfg_links( + ctx, + acquire, + reusable_sync, + ready_sync, + half_zero, + commit, + tail, + ) +} + +/// 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)?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::attributes::{ + CutePipelineStateAttr, CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, + CuteTensorLayoutAttr, + }; + use crate::layout::{ComposedLayout, OffsetUnit}; + use dialect_mir::ops::{MirReturnOp, MirUndefOp}; + use dialect_mir::types::{MirPtrType, address_space}; + use pliron::basic_block::BasicBlock; + use pliron::builtin::attributes::TypeAttr; + use pliron::builtin::op_interfaces::{OperandSegmentInterface, SymbolOpInterface}; + use pliron::builtin::ops::ModuleOp; + use pliron::builtin::types::{FP32Type, FunctionType, IntegerType, Signedness}; + + 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) + } + + 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) + } + + 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() + } + + 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) + } + + 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); + } + + 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) + } + + 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 + } + + 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); + } + + 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); + } + + fn append_complete_scheduler_function( + ctx: &mut Context, + module_block: Ptr, + name: &str, + grid: CuteTileGridAttr, + misroute_work_edge: bool, + ) { + let function_type = FunctionType::get(ctx, vec![], vec![]); + let function = Operation::new( + ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function_op = MirFuncOp::new(ctx, function, TypeAttr::new(function_type.into())); + function_op.set_symbol_name(ctx, name.try_into().unwrap()); + function.insert_at_back(module_block, ctx); + let region = function.deref(ctx).get_region(0); + let u64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + let entry = BasicBlock::new(ctx, None, vec![]); + entry.insert_at_back(region, ctx); + let header = BasicBlock::new(ctx, None, vec![u64_ty, u64_ty]); + header.insert_at_back(region, ctx); + let body = BasicBlock::new(ctx, None, vec![]); + body.insert_at_back(region, ctx); + let exit = BasicBlock::new(ctx, None, vec![]); + exit.insert_at_back(region, ctx); + + let start = CuteSchedulerNew1dOp::new(ctx, grid).get_operation(); + start.insert_at_back(entry, ctx); + let start = CuteSchedulerNew1dOp::wrap(start); + goto( + ctx, + entry, + header, + vec![start.current(ctx), start.stride(ctx)], + ); + + let current = header.deref(ctx).get_argument(0); + let stride = header.deref(ctx).get_argument(1); + let has_work = CuteSchedulerHasWorkOp::new(ctx, current, grid).get_operation(); + has_work.insert_at_back(header, ctx); + let not = Operation::new( + ctx, + MirNotOp::get_concrete_op_info(), + vec![ + CuteSchedulerHasWorkOp::wrap(has_work) + .has_work(ctx) + .get_type(ctx), + ], + vec![CuteSchedulerHasWorkOp::wrap(has_work).has_work(ctx)], + vec![], + 0, + ); + not.insert_at_back(header, ctx); + let finished = not.deref(ctx).get_result(0); + let (true_target, false_target) = if misroute_work_edge { + (body, exit) + } else { + (exit, body) + }; + cond_branch(ctx, header, finished, true_target, false_target); + + let selected = CuteSchedulerCurrentOp::new(ctx, current, grid).get_operation(); + selected.insert_at_back(body, ctx); + CuteWorkTileCoordinatesOp::new(ctx, CuteSchedulerCurrentOp::wrap(selected).work_tile(ctx)) + .get_operation() + .insert_at_back(body, ctx); + let advance = CuteSchedulerAdvanceOp::new(ctx, current, stride, grid).get_operation(); + advance.insert_at_back(body, ctx); + goto( + ctx, + body, + header, + vec![CuteSchedulerAdvanceOp::wrap(advance).next(ctx), stride], + ); + + Operation::new( + ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ) + .insert_at_back(exit, ctx); + } + + #[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}"); + } + + #[test] + fn accepts_independent_scheduler_stories_in_two_functions_without_mutation() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + append_complete_scheduler_function( + &mut ctx, + block, + "first", + CuteTileGridAttr::new(4, 4, 1), + false, + ); + append_complete_scheduler_function( + &mut ctx, + block, + "second", + CuteTileGridAttr::new(5, 4, 1), + false, + ); + let mut before = Vec::new(); + collect_ops(&ctx, module, &mut before); + + verify_cute_semantics(&ctx, module).unwrap(); + + let mut after = Vec::new(); + collect_ops(&ctx, module, &mut after); + assert_eq!(after, before); + } + + #[test] + fn rejects_a_scheduler_body_on_the_exit_edge_even_when_preorder_looks_ordered() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + append_complete_scheduler_function( + &mut ctx, + block, + "misrouted", + CuteTileGridAttr::new(4, 4, 1), + true, + ); + + let error = verify_cute_semantics(&ctx, module).unwrap_err().to_string(); + assert!(error.contains("scheduler CFG"), "{error}"); + } + + #[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 + ); + } + + #[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}"); + } + + #[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::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::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::local_cell(&ctx, escalated_alias).is_none()); + assert!(!local_cell_is_closed(&ctx, pointer)); + } + + #[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 + )); + } + + #[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}"); + } + + #[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}"); + } + + #[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}" + ); + } + + #[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}"); + } + + #[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}"); + } + + #[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}"); + } + + #[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/cuteir/examples/blockscale_gemm_cute/Cargo.lock b/cuteir/examples/blockscale_gemm_cute/Cargo.lock new file mode 100644 index 0000000000..6a848bbf7f --- /dev/null +++ b/cuteir/examples/blockscale_gemm_cute/Cargo.lock @@ -0,0 +1,567 @@ +# 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.71.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn 2.0.119", +] + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[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 = [ + "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.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.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.2.1" +dependencies = [ + "bindgen", +] + +[[package]] +name = "cuda-core" +version = "0.2.1" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-macros", + "half", + "oxide-artifacts", +] + +[[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", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn 2.0.119", +] + +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "cute-rs" +version = "0.1.0" +dependencies = [ + "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.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[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 = "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.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + +[[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 = [ + "libloading", + "thiserror", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[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 = "oxide-artifacts" +version = "0.2.1" +dependencies = [ + "object", +] + +[[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 = "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 = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[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.3", +] + +[[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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[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 = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "zerocopy" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556764e583adb45a9f8d413c2a147fa7e8d821e48e12b14fd560b607998b75eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ab42fc20575779bd240faa45f94a74256f755c0fa9e89f0ede20d91d0cdfc1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] diff --git a/cuteir/examples/blockscale_gemm_cute/Cargo.toml b/cuteir/examples/blockscale_gemm_cute/Cargo.toml new file mode 100644 index 0000000000..5a1aa09166 --- /dev/null +++ b/cuteir/examples/blockscale_gemm_cute/Cargo.toml @@ -0,0 +1,16 @@ +[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] +cute-rs = { path = "../../cute-rs", features = ["host"] } +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../crates/cuda-core" } diff --git a/cuteir/examples/blockscale_gemm_cute/README.md b/cuteir/examples/blockscale_gemm_cute/README.md new file mode 100644 index 0000000000..1705edb4eb --- /dev/null +++ b/cuteir/examples/blockscale_gemm_cute/README.md @@ -0,0 +1,122 @@ +# Block-scaled MXFP4 GEMM + +GEMM means “matrix times matrix”: + +```text +C[M,N] = A[M,K] × B[K,N] + +M = 2048 N = 2048 K = 1024 +input values = packed E2M1 (FP4) +input scale = one UE8M0 value per 32 K values +accumulator = f32 +output = f16 +``` + +The shared scale moves each group of small FP4 values into a useful numerical +range. The scale is applied as part of the matrix multiply. + +The output is split into `128 × 128` tiles. A thread block, also called a CTA, +works on one output tile at a time. + +The four hardware terms used below are: + +- **TMA**, or Tensor Memory Accelerator, is a hardware engine that copies a + whole multidimensional tile; +- **shared memory** is fast storage shared by the threads in one CTA; +- **MMA** is the tensor-core matrix multiply-and-accumulate operation; +- an **epilogue** converts and arranges completed accumulators for output. + +## The Tensor flow + +This example follows the same shape as a large CuTe kernel: + +```text +global A Tensor ─┐ +global B Tensor ─┼── TMA tile copies ──► shared-memory Tensors +A/B scale Tensors┘ │ + ▼ + per-warp tiled views + │ + load A/B/scale fragments + │ + ▼ + block-scaled MMA + │ + ▼ + f32 accumulator tile + │ + tiled f16 epilogue + │ + ▼ + shared output Tensor + │ + TMA tile stores + │ + ▼ + global C Tensor +``` + +## How work moves through the kernel + +The kernel uses three shared-memory slots: + +```text +slot 0 slot 1 slot 2 +load here → compute here → waiting here → repeat +``` + +One producer warp issues the TMA loads. Eight compute warps consume completed +slots and update their accumulator tiles. While the compute warps use one slot, +TMA can fill another. + +After a CTA finishes an output tile, a small scheduler may give it another +tile: + +```text +physical CTA 0: logical tile 0, then tile 0 + number_of_CTAs, ... +physical CTA 1: logical tile 1, then tile 1 + number_of_CTAs, ... +``` + +This is called persistent scheduling. It changes which CTA owns a tile; it +does not change the matrix calculation. + +The Rust types keep the important agreements together. For example, a +`TmaDesc` built on the host must name the same shared-memory layout +as the device copy that consumes it. Matrix A and matrix B also carry different +operand-role types, so they cannot be silently swapped in a tiled MMA call. + +## Device requirement + +The packed MXFP4 tensor-core instruction used here requires an SM 12.x +Blackwell target (`sm_120` or `sm_121`). Build the example for the accelerated +target `sm_120a`. + +On another GPU, the program skips device execution. Numerical GEMM validation +requires a compatible SM 12.x GPU. + +## Run it + +From the repository root: + +```bash +cargo oxide run blockscale_gemm_cute --arch sm_120a +``` + +The program builds deterministic packed inputs, computes a host reference, +runs the GPU kernel, and compares every output bit pattern. The host reference +for the fixed `2048 × 2048 × 1024` problem is intentionally substantial, so a +checked run is not instant. + +A successful device run confirms all 4,194,304 `f16` output bit patterns. It +then performs 100 warmup launches and measures 31 individual launches with +CUDA events. The report includes the median, p10, and p90 device time and the +median-derived TFLOP/s for `2 × M × N × K` floating-point operations. + +This is the full semantic acceptance example: + +```text +Tensor layouts → tiled copies → pipeline → tiled MMA → output Tensor +``` + +The backend receives the same high-level module rather than a second kernel +description. diff --git a/cuteir/examples/blockscale_gemm_cute/src/main.rs b/cuteir/examples/blockscale_gemm_cute/src/main.rs new file mode 100644 index 0000000000..3f90dd3daf --- /dev/null +++ b/cuteir/examples/blockscale_gemm_cute/src/main.rs @@ -0,0 +1,684 @@ +/* + * 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.accumulate_k64(...) +//! typed epilogue: AccC -> f16x2 -> stmatrix -> shared +//! warp 0, lane 0: two 128x64 TMA stores -> global C +//! } +//! producer/consumer slot + phase survive between output tiles +//! ``` +//! +//! `TmaDesc` 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 cuda_core::{CudaContext, LaunchConfig1D}; +use cuda_device::barrier::Barrier; +use cuda_device::{ + DynamicSharedArray, cuda_module, kernel, launch_bounds, launch_contract, thread, warp, +}; +use cute_rs::tma::{TmaEncodeOptions, make_tma_desc_2d_with_options}; +use cute_rs::{ + C, Composed, Consumer, L, Mkl, Mxf4AccumulatorTile2x8, Mxf4E2M1, Mxfp4TiledMma, Nkl, + PipelineState, Producer, RowMajor, SharedScaleAtom, SharedTensor, Sm1xxBlockScaleKMajor, + Sm120Epilogue128x128, Sm120EpilogueHalfLayout, Sm120ScaleAtom, SmemTile, + StaticPersistentTileScheduler, Swizzle, T2, TmaDesc, TmaLoadPipeline, TmaStorePipeline, + copy_tma_2d, copy_tma_s2g_2d, +}; + +const M: usize = 2048; +const N: usize = 2048; +const K: usize = 1024; // logical e2m1 elements; four packed into each f16 carrier +const KB: usize = K / 32; // K-blocks: one ue8m0 scale per 32 elements +const COMPUTE_WARPS: usize = 8; // 4 along M, 2 along N +const PRODUCER_WARP: usize = COMPUTE_WARPS; +const THREADS: usize = (COMPUTE_WARPS + 1) * 32; +const BLOCK_M: usize = 128; +const BLOCK_N: usize = 128; +const BATCHES: usize = 1; +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. +type TileP = L, C<32>>, T2, C<1>>>; +type SmemP = Composed, 0, TileP>; +type ByteSmem = Composed, 0, RowMajor<128, 64>>; +type ScaleTmaSmem = RowMajor<1, 256>; +type SharedA = SharedTensor; +type SharedB = SharedTensor; + +const K_STAGES: usize = KB / 4; +const DATA_STAGE_BYTES: usize = BLOCK_M * 64; +const SCALE_STAGE_BYTES: usize = Sm120ScaleAtom::BYTES; +const TMA_TX_BYTES: u32 = (2 * DATA_STAGE_BYTES + 2 * SCALE_STAGE_BYTES) as u32; +const PIPELINE_STAGES: usize = 3; +type MainloopPipeline = TmaLoadPipeline; +type EpiloguePipeline = TmaStorePipeline<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. +const SMEM_A_OFFSET: usize = 0; +const SMEM_B_OFFSET: usize = SMEM_A_OFFSET + PIPELINE_STAGES * DATA_STAGE_BYTES; +const SMEM_SFA_OFFSET: usize = SMEM_B_OFFSET + PIPELINE_STAGES * DATA_STAGE_BYTES; +const SMEM_SFB_OFFSET: usize = SMEM_SFA_OFFSET + PIPELINE_STAGES * SCALE_STAGE_BYTES; +const PIPELINE_OFFSET: usize = SMEM_SFB_OFFSET + PIPELINE_STAGES * SCALE_STAGE_BYTES; +const MAINLOOP_SMEM_BYTES: usize = PIPELINE_OFFSET + MainloopPipeline::STORAGE_BYTES; +const EPILOGUE_OFFSET: usize = + MAINLOOP_SMEM_BYTES.next_multiple_of(Sm120Epilogue128x128::ALIGNMENT); +const DYNAMIC_SMEM_BYTES: usize = EPILOGUE_OFFSET + Sm120Epilogue128x128::BYTES; +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); +}; + +#[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 TmaDesc, + b_tma: *const TmaDesc, + sfa_tma: *const TmaDesc, + sfb_tma: *const TmaDesc, + out_tma: *const TmaDesc, + ) { + 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 = Mxfp4TiledMma::::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.current_tile(); + let (_, m_tile, n_tile, batch) = work_tile.coordinates(); + 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.slot(); + 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.accumulate_k64( + &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(); + } + 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.current_tile(); + let (_, m_tile, n_tile, batch) = work_tile.coordinates(); + 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.slot(); + 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(); + } + // 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]; + +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 +} + +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 +} + +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 +} + +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 + .chunks_exact(4) + .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.offset(0, row, group)] = dense[row * KB + group]; + } + } + packed + .chunks_exact(2) + .map(|pair| u16::from_le_bytes([pair[0], pair[1]])) + .collect() +} + +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] +} + +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: TmaDesc = make_tma_desc_2d_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: TmaDesc = make_tma_desc_2d_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: TmaDesc = make_tma_desc_2d_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: TmaDesc = make_tma_desc_2d_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 TmaDesc; + let b_desc_ptr = b_desc_dev.cu_deviceptr() as *const TmaDesc; + let sfa_desc_ptr = sfa_desc_dev.cu_deviceptr() as *const TmaDesc; + let sfb_desc_ptr = sfb_desc_dev.cu_deviceptr() as *const TmaDesc; + + let out = cuda_core::DeviceBuffer::::zeroed(&stream, M * N).unwrap(); + let out_desc: TmaDesc = make_tma_desc_2d_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 TmaDesc; + // 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::resident_grid_size(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. + const WARMUP: usize = 100; + 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/cuteir/examples/elementwise_cute/Cargo.lock b/cuteir/examples/elementwise_cute/Cargo.lock new file mode 100644 index 0000000000..f8ab63ab9f --- /dev/null +++ b/cuteir/examples/elementwise_cute/Cargo.lock @@ -0,0 +1,566 @@ +# 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.71.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn 2.0.119", +] + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[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.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.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.2.1" +dependencies = [ + "bindgen", +] + +[[package]] +name = "cuda-core" +version = "0.2.1" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-macros", + "half", + "oxide-artifacts", +] + +[[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", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn 2.0.119", +] + +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "cute-rs" +version = "0.1.0" +dependencies = [ + "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.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" + +[[package]] +name = "elementwise_cute" +version = "0.1.0" +dependencies = [ + "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 = "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 = "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.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + +[[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 = [ + "libloading", + "thiserror", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[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 = "oxide-artifacts" +version = "0.2.1" +dependencies = [ + "object", +] + +[[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 = "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 = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[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.3", +] + +[[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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[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 = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "zerocopy" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556764e583adb45a9f8d413c2a147fa7e8d821e48e12b14fd560b607998b75eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ab42fc20575779bd240faa45f94a74256f755c0fa9e89f0ede20d91d0cdfc1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] diff --git a/cuteir/examples/elementwise_cute/Cargo.toml b/cuteir/examples/elementwise_cute/Cargo.toml new file mode 100644 index 0000000000..3abd0aee2a --- /dev/null +++ b/cuteir/examples/elementwise_cute/Cargo.toml @@ -0,0 +1,16 @@ +[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] +cute-rs = { path = "../../cute-rs" } +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../crates/cuda-core" } diff --git a/cuteir/examples/elementwise_cute/README.md b/cuteir/examples/elementwise_cute/README.md new file mode 100644 index 0000000000..c6fa4df525 --- /dev/null +++ b/cuteir/examples/elementwise_cute/README.md @@ -0,0 +1,93 @@ +# 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); + +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 normal CUDA Oxide toolchain and an NVIDIA GPU supported by the +project. From the repository root: + +```bash +cargo oxide run elementwise_cute +``` + +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/cuteir/examples/elementwise_cute/src/main.rs b/cuteir/examples/elementwise_cute/src/main.rs new file mode 100644 index 0000000000..038d5d9c3a --- /dev/null +++ b/cuteir/examples/elementwise_cute/src/main.rs @@ -0,0 +1,167 @@ +/* + * 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 cuda_core::{CudaContext, DeviceBuffer, LaunchConfig}; +use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; + +/// Deliberately not divisible by the tiles: f32 tail of 3, f16 tail of 5. +const N_F32: usize = 1003; +const N_F16: usize = 2005; +const TILE_F32: usize = 4; // 16 bytes +const TILE_F16: usize = 8; // 16 bytes + +#[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 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 { + // Ragged tail: plain Rust scalar loop, boundary thread only. + let mut k = t_a.base(); + while k < a.len() { + unsafe { t_c.store_linear(k, a[k] + b[k]) }; + k += 1; + } + } + } + + #[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_linear(k, a[k] + b[k]) }; + k += 1; + } + } + } +} + +fn ceil_div(a: usize, b: usize) -> usize { + a.div_ceil(b) +} + +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/cuteir/examples/nvfp4_gemv_cute/Cargo.lock b/cuteir/examples/nvfp4_gemv_cute/Cargo.lock new file mode 100644 index 0000000000..9394caaae8 --- /dev/null +++ b/cuteir/examples/nvfp4_gemv_cute/Cargo.lock @@ -0,0 +1,566 @@ +# 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.71.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn 2.0.119", +] + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[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.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.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.2.1" +dependencies = [ + "bindgen", +] + +[[package]] +name = "cuda-core" +version = "0.2.1" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-macros", + "half", + "oxide-artifacts", +] + +[[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", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn 2.0.119", +] + +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "cute-rs" +version = "0.1.0" +dependencies = [ + "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.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[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 = "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.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + +[[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 = [ + "libloading", + "thiserror", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[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 = [ + "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 = "oxide-artifacts" +version = "0.2.1" +dependencies = [ + "object", +] + +[[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 = "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 = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[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.3", +] + +[[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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[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 = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "zerocopy" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556764e583adb45a9f8d413c2a147fa7e8d821e48e12b14fd560b607998b75eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ab42fc20575779bd240faa45f94a74256f755c0fa9e89f0ede20d91d0cdfc1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] diff --git a/cuteir/examples/nvfp4_gemv_cute/Cargo.toml b/cuteir/examples/nvfp4_gemv_cute/Cargo.toml new file mode 100644 index 0000000000..39f3d26dac --- /dev/null +++ b/cuteir/examples/nvfp4_gemv_cute/Cargo.toml @@ -0,0 +1,17 @@ +[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] +cute-rs = { path = "../../cute-rs" } +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../crates/cuda-core" } diff --git a/cuteir/examples/nvfp4_gemv_cute/README.md b/cuteir/examples/nvfp4_gemv_cute/README.md new file mode 100644 index 0000000000..3f333336b9 --- /dev/null +++ b/cuteir/examples/nvfp4_gemv_cute/README.md @@ -0,0 +1,103 @@ +# 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, E2M1, UE8M0, 16, KMajor>; +type VectorB<'a> = + BlockScaledTensor<'a, E2M1, UE8M0, 16, KMajor>; +``` + +They record: + +- the packed value format (`E2M1`); +- the scale format (`UE8M0`); +- 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`. + +## Run it + +From the repository root: + +```bash +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. diff --git a/cuteir/examples/nvfp4_gemv_cute/src/main.rs b/cuteir/examples/nvfp4_gemv_cute/src/main.rs new file mode 100644 index 0000000000..d4b921ccbd --- /dev/null +++ b/cuteir/examples/nvfp4_gemv_cute/src/main.rs @@ -0,0 +1,423 @@ +/* + * 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 cuda_core::{CudaContext, DeviceBuffer, LaunchConfig}; +use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; +use cute_rs::{BlockScaledTensor, E2M1, KMajor, Mkl, Nkl, Sm1xxBlockScaleKMajor, UE8M0}; + +const CTA_M: usize = 128; +const K_TILE: usize = 64; +const SF_VEC: usize = 16; +const DEFAULT_M: usize = 512; +const DEFAULT_K: usize = 256; +const DEFAULT_L: usize = 1; + +type MatrixA<'a> = BlockScaledTensor<'a, E2M1, UE8M0, SF_VEC, KMajor>; + +type VectorB<'a> = BlockScaledTensor<'a, E2M1, UE8M0, SF_VEC, KMajor>; + +#[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); + } + } +} + +#[derive(Clone, Copy, Debug)] +struct Args { + m: usize, + k: usize, + l: usize, + device: usize, +} + +fn usage() -> ! { + println!("Usage: nvfp4_gemv_cute [--m M] [--k K] [--l L] [--device ORDINAL]"); + std::process::exit(0); +} + +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 +} + +fn ceil_div(a: usize, b: usize) -> usize { + a.div_ceil(b) +} + +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")) +} + +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"); +} + +fn e2m1_to_f32(nibble: u8) -> f32 { + 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 + } +} + +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) + } +} + +fn a_nibble(batch: usize, row: usize, k: usize) -> u8 { + (((batch % 16) * 5 + (row % 16) * 13 + (k % 16) * 7 + 3) % 16) as u8 +} + +fn b_nibble(batch: usize, k: usize) -> u8 { + (((batch % 16) * 3 + (k % 16) * 11 + 1) % 16) as u8 +} + +fn sfa_byte(batch: usize, row: usize, k_group: usize) -> u8 { + 126 + ((batch % 3 + row % 3 + k_group % 3) % 3) as u8 +} + +fn sfb_byte(batch: usize, k_group: usize) -> u8 { + 126 + ((batch % 3 + 2 * (k_group % 3)) % 3) as u8 +} + +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 +} + +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 +} + +fn scale_storage_len(l: usize, mn: usize, k: usize) -> usize { + Sm1xxBlockScaleKMajor::::new(mn, k).storage_len(l) +} + +fn scale_index(batch: usize, row: usize, k_group: usize, mn: usize, k: usize) -> usize { + Sm1xxBlockScaleKMajor::::new(mn, k).offset(batch, row, k_group) +} + +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) +} + +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 +} + +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) + ); +} + +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) + }) +} + +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"); +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::BTreeSet; + + #[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); + } + + #[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()); + } + + #[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); + } + + #[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/dependency-licenses.csv b/dependency-licenses.csv index 79325bd645..afd2e56ef4 100644 --- a/dependency-licenses.csv +++ b/dependency-licenses.csv @@ -36,8 +36,11 @@ cuda-intrinsics-gen,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace) cuda-macros,0.2.1,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,"Procedural macros for CUDA kernel development (#[kernel], #[device], etc.)", cuda-oxide-codegen,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,Rustc-independent cuda-oxide PTX backend, cuda-target-spec,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,Shared CUDA target parsing and recorded LLVM PTX-floor policy, +cute-layout,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,Coordinate-to-memory mapping and validation for CuTe-style GPU tiles, +cute-rs,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,"CuTe-style device library: Tensor, Layout types, partitioning, copy/mma atoms", dashmap,6.2.1,MIT,Yes,registry+https://github.com/rust-lang/crates.io-index,3rd Party (crates.io),https://github.com/xacrimon/dashmap,Blazing fast concurrent HashMap for Rust.,Pulled by the shared cuda-async (cutile-rs). device-only-kernels,0.0.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,"INTERNAL: fixture for scripts/check-device-only-build.sh, proving the macro expansion resolves with no cuda-host in the graph. Not published.",Declares its own [workspace] so it resolves independently of the root lock. +dialect-cute,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,Semantic CuTe dialect for retained tensor layout and scheduling operations, dialect-iket,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,Semantic IKET dialect for in-kernel event tracing, dialect-mir,0.2.1,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,MIR dialect for pliron MLIR framework (Rust MIR representation), dialect-nvvm,0.2.1,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,NVVM dialect for pliron MLIR framework (NVIDIA GPU intrinsics), diff --git a/scripts/check-dependency-licenses.sh b/scripts/check-dependency-licenses.sh index 6c59652c41..59ca869efc 100755 --- a/scripts/check-dependency-licenses.sh +++ b/scripts/check-dependency-licenses.sh @@ -137,7 +137,10 @@ VENDORED_WORKSPACE_ROOTS=( # Cargo.toml -> Cargo workspace root # |-- named first-party or vendored root # `-- example root with a tracked Cargo.lock -EXAMPLES_ROOT=crates/rustc-codegen-cuda/examples +EXAMPLE_ROOTS=( + crates/rustc-codegen-cuda/examples + cuteir/examples +) all_workspace_roots="$( git ls-files -z -- '*Cargo.toml' | while IFS= read -r -d '' manifest; do @@ -162,7 +165,7 @@ fi on_disk_roots="$( printf '%s\n' "${all_workspace_roots}" | - grep -v "^${EXAMPLES_ROOT}/" | + grep -v -e "^${EXAMPLE_ROOTS[0]}/" -e "^${EXAMPLE_ROOTS[1]}/" | LC_ALL=C sort )" named_roots="$( @@ -198,7 +201,7 @@ fi # the two sets before reading any package names so a new nested workspace # cannot disappear merely because it has not committed a lock yet. mapfile -d '' -t EXAMPLE_LOCKFILES < <( - git ls-files -z -- "${EXAMPLES_ROOT}/**/Cargo.lock" + git ls-files -z -- "${EXAMPLE_ROOTS[0]}/**/Cargo.lock" "${EXAMPLE_ROOTS[1]}/**/Cargo.lock" ) if [[ ${#EXAMPLE_LOCKFILES[@]} -lt 20 ]]; then echo "error: found only ${#EXAMPLE_LOCKFILES[@]} tracked example lock files;" >&2 @@ -207,7 +210,7 @@ if [[ ${#EXAMPLE_LOCKFILES[@]} -lt 20 ]]; then fi example_roots="$( printf '%s\n' "${all_workspace_roots}" | - grep "^${EXAMPLES_ROOT}/" | + grep -e "^${EXAMPLE_ROOTS[0]}/" -e "^${EXAMPLE_ROOTS[1]}/" | LC_ALL=C sort )" example_lock_roots="$( @@ -258,7 +261,7 @@ echo "OK: ${CSV} records all $(printf '%s\n' "${required}" | grep -c .) declared # --------------------------------------------------------------------------- # Second half: the example workspaces. # -# Every example under crates/rustc-codegen-cuda/examples/ sets its own +# Every example under the ordinary and CuTe example roots sets its own # [workspace], so neither `cargo deny check` nor the check above resolves any # of them -- both stop at the root workspace boundary. Most examples declare # only path dependencies on first-party crates and so bring nothing new, but a @@ -316,11 +319,14 @@ locks = sorted(sys.argv[separator + 1:]) if len(locks) < 20: sys.exit("parse self-test failed: found %d example lock files" % len(locks)) -examples_root = "crates/rustc-codegen-cuda/examples" +example_roots = ("crates/rustc-codegen-cuda/examples", "cuteir/examples") def example_of(lock): """Top-level example directory a lockfile belongs to, however deep it sits.""" - return os.path.relpath(lock, examples_root).split(os.sep)[0] + for root in example_roots: + if lock.startswith(root + os.sep): + return os.path.relpath(lock, root).split(os.sep)[0] + sys.exit("lockfile outside every configured example root: " + lock) present = {example_of(lock) for lock in locks} unknown = sorted(exempt - present) diff --git a/scripts/check-example-license-policy.sh b/scripts/check-example-license-policy.sh index 11203108f3..0168e0f7ee 100644 --- a/scripts/check-example-license-policy.sh +++ b/scripts/check-example-license-policy.sh @@ -4,8 +4,8 @@ # Enforce deny.toml over the example workspaces, which `cargo deny check` does # not reach. # -# `cargo deny check` resolves the root workspace. Every example under -# crates/rustc-codegen-cuda/examples/ sets its own `[workspace]`, so the root +# `cargo deny check` resolves the root workspace. Every example under the +# roots below sets its own `[workspace]`, so the root # run stops at that boundary and the license, source and ban policies never see # any crate an example pulls on its own. #664 closed the *inventory* half of # this (dependency-licenses.csv now records those crates); this closes the @@ -36,7 +36,10 @@ export LC_ALL=C cd "$(dirname "$0")/.." -EXAMPLES_ROOT=crates/rustc-codegen-cuda/examples +EXAMPLE_ROOTS=( + crates/rustc-codegen-cuda/examples + cuteir/examples +) # Examples whose dependencies deliberately cannot satisfy deny.toml. # @@ -65,7 +68,11 @@ command -v cargo-deny >/dev/null 2>&1 || { representatives="$(python3 -c ' import glob, os, re, sys -examples_root, *exempt = sys.argv[1:] +separator = sys.argv.index("--roots") +exempt = sys.argv[1:separator] +example_roots = sys.argv[separator + 1:] +if not example_roots: + sys.exit("parse self-test failed: no example roots configured") def third_party(lock): """(name, version, source) for every locked package that has a source. @@ -92,8 +99,15 @@ def third_party(lock): found.add((name.group(1), version.group(1), source.group(1))) return found, parsed -locks = sorted(glob.glob(os.path.join(examples_root, "**", "Cargo.lock"), recursive=True)) -on_disk = {os.path.relpath(lock, examples_root).split(os.sep)[0] for lock in locks} +locks = sorted( + (root, lock) + for root in example_roots + for lock in glob.glob(os.path.join(root, "**", "Cargo.lock"), recursive=True) +) +on_disk = { + os.path.relpath(lock, root).split(os.sep)[0] + for root, lock in locks +} # Cross-check the glob against an independent enumeration rather than a fixed # floor. Every example directory that carries a Cargo.toml carries a @@ -104,8 +118,10 @@ on_disk = {os.path.relpath(lock, examples_root).split(os.sep)[0] for lock in loc # not to fail on ordinary growth. The floor this replaces was 20. expected = { name - for name in os.listdir(examples_root) - if os.path.exists(os.path.join(examples_root, name, "Cargo.toml")) + for root in example_roots + if os.path.isdir(root) + for name in os.listdir(root) + if os.path.exists(os.path.join(root, name, "Cargo.toml")) } missed = sorted(expected - on_disk) if missed: @@ -121,8 +137,8 @@ if unknown: groups = {} total_parsed = 0 -for lock in locks: - example = os.path.relpath(lock, examples_root).split(os.sep)[0] +for root, lock in locks: + example = os.path.relpath(lock, root).split(os.sep)[0] crates, parsed = third_party(lock) total_parsed += parsed # Every example reaches crates.io through cuda-core/cuda-device/cuda-host, @@ -149,10 +165,10 @@ if total_parsed < 10 * len(locks): for manifest in sorted(groups.values()): print(manifest) -' "${EXAMPLES_ROOT}" "${POLICY_EXEMPT_EXAMPLES[@]}")" +' "${POLICY_EXEMPT_EXAMPLES[@]}" --roots "${EXAMPLE_ROOTS[@]}")" total="$(printf '%s\n' "${representatives}" | grep -c .)" -locks="$(find "${EXAMPLES_ROOT}" -name Cargo.lock | grep -c .)" +locks="$(find "${EXAMPLE_ROOTS[@]}" -name Cargo.lock | grep -c .)" echo "Checking deny.toml over ${total} representative example workspaces" \ "across ${locks} example lock files." @@ -165,7 +181,7 @@ echo "Checking deny.toml over ${total} representative example workspaces" \ failed=() for manifest in ${representatives}; do if ! cargo deny --manifest-path "${manifest}" --locked check 2>&1 | - sed "s|^| [$(dirname "${manifest#"${EXAMPLES_ROOT}/"}")] |"; then + sed "s|^| [$(dirname "${manifest}")] |"; then failed+=("${manifest}") fi done From 4341206cfd3d435a88b726f26f6d63dc83d63e55 Mon Sep 17 00:00:00 2001 From: nihalp Date: Mon, 17 Aug 2026 23:49:05 +0530 Subject: [PATCH 02/99] feat(cuteir): add CUTLASS translation backend --- .gitignore | 1 + Cargo.lock | 31 + Cargo.toml | 5 + README.md | 4 +- .../cargo-oxide/src/commands/codegen_env.rs | 56 + .../cargo-oxide/src/commands/fingerprint.rs | 4 +- .../cargo-oxide/src/commands/passthrough.rs | 77 +- crates/cargo-oxide/src/commands/tests.rs | 129 + crates/cargo-oxide/src/cutlass_toolchain.rs | 1033 ++++++ crates/cargo-oxide/src/main.rs | 51 + crates/cuda-oxide-codegen/Cargo.toml | 6 + crates/cuda-oxide-codegen/src/api.rs | 4 + .../cuda-oxide-codegen/src/cutlass_backend.rs | 1187 ++++++ .../src/cutlass_compiler.rs | 896 +++++ crates/cuda-oxide-codegen/src/error.rs | 3 + crates/cuda-oxide-codegen/src/lib.rs | 4 + crates/cuda-oxide-codegen/src/pipeline.rs | 455 +++ crates/cuda-oxide-mlir-export/Cargo.toml | 23 + crates/cuda-oxide-mlir-export/README.md | 157 + crates/cuda-oxide-mlir-export/src/cute.rs | 1000 ++++++ .../cuda-oxide-mlir-export/src/cute_gemm.rs | 3190 +++++++++++++++++ .../cuda-oxide-mlir-export/src/cute_gemv.rs | 1243 +++++++ crates/cuda-oxide-mlir-export/src/lib.rs | 27 + crates/cuda-oxide-mlir-export/src/mir_core.rs | 1273 +++++++ .../cuda-oxide-mlir-export/src/mir_memory.rs | 1308 +++++++ .../cuda-oxide-mlir-export/src/nvvm_sregs.rs | 256 ++ crates/cuda-oxide-mlir-export/src/packs.rs | 300 ++ crates/cuda-oxide-mlir-export/src/profile.rs | 273 ++ .../tests/cute_elementwise.rs | 439 +++ .../cuda-oxide-mlir-export/tests/cute_gemv.rs | 260 ++ crates/mir-importer/src/lib.rs | 1 + crates/mir-importer/src/pipeline.rs | 44 +- crates/pliron-mlir-export/Cargo.toml | 16 + crates/pliron-mlir-export/README.md | 20 + crates/pliron-mlir-export/src/ast.rs | 206 ++ crates/pliron-mlir-export/src/error.rs | 76 + crates/pliron-mlir-export/src/lib.rs | 42 + crates/pliron-mlir-export/src/registry.rs | 253 ++ crates/pliron-mlir-export/src/render.rs | 472 +++ crates/pliron-mlir-export/src/translate.rs | 839 +++++ crates/rustc-codegen-cuda/Cargo.lock | 29 + crates/rustc-codegen-cuda/README.md | 27 +- .../rustc-codegen-cuda/src/device_codegen.rs | 12 + crates/rustc-codegen-cuda/src/lib.rs | 157 +- .../getting-started/installation.md | 24 + cuteir/README.md | 5 + cuteir/docs/translation.md | 77 + .../examples/blockscale_gemm_cute/Cargo.lock | 5 + cuteir/examples/elementwise_cute/Cargo.lock | 5 + cuteir/examples/nvfp4_gemv_cute/Cargo.lock | 5 + dependency-licenses.csv | 2 + 51 files changed, 16002 insertions(+), 10 deletions(-) create mode 100644 crates/cargo-oxide/src/cutlass_toolchain.rs create mode 100644 crates/cuda-oxide-codegen/src/cutlass_backend.rs create mode 100644 crates/cuda-oxide-codegen/src/cutlass_compiler.rs create mode 100644 crates/cuda-oxide-mlir-export/Cargo.toml create mode 100644 crates/cuda-oxide-mlir-export/README.md create mode 100644 crates/cuda-oxide-mlir-export/src/cute.rs create mode 100644 crates/cuda-oxide-mlir-export/src/cute_gemm.rs create mode 100644 crates/cuda-oxide-mlir-export/src/cute_gemv.rs create mode 100644 crates/cuda-oxide-mlir-export/src/lib.rs create mode 100644 crates/cuda-oxide-mlir-export/src/mir_core.rs create mode 100644 crates/cuda-oxide-mlir-export/src/mir_memory.rs create mode 100644 crates/cuda-oxide-mlir-export/src/nvvm_sregs.rs create mode 100644 crates/cuda-oxide-mlir-export/src/packs.rs create mode 100644 crates/cuda-oxide-mlir-export/src/profile.rs create mode 100644 crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs create mode 100644 crates/cuda-oxide-mlir-export/tests/cute_gemv.rs create mode 100644 crates/pliron-mlir-export/Cargo.toml create mode 100644 crates/pliron-mlir-export/README.md create mode 100644 crates/pliron-mlir-export/src/ast.rs create mode 100644 crates/pliron-mlir-export/src/error.rs create mode 100644 crates/pliron-mlir-export/src/lib.rs create mode 100644 crates/pliron-mlir-export/src/registry.rs create mode 100644 crates/pliron-mlir-export/src/render.rs create mode 100644 crates/pliron-mlir-export/src/translate.rs create mode 100644 cuteir/docs/translation.md diff --git a/.gitignore b/.gitignore index d87b0336a5..ef920ba635 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,7 @@ target/ .ltoir_arch *.ll *.bc +*.mlir # Generated intrinsic probes are checked-in fixtures. !intrinsics/probes/*.ll diff --git a/Cargo.lock b/Cargo.lock index f9ff24019f..0379b70843 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -483,24 +483,46 @@ 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", + "pliron-mlir-export", "ptx-parse", "rustc-demangle", + "serde_json", + "sha2", "tempfile", "thiserror 2.0.18", ] +[[package]] +name = "cuda-oxide-mlir-export" +version = "0.1.0" +dependencies = [ + "dialect-cute", + "dialect-mir", + "dialect-nvvm", + "llvm-export", + "pliron", + "pliron-derive", + "pliron-mlir-export", + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -1208,6 +1230,15 @@ dependencies = [ "thiserror 2.0.18", ] +[[package]] +name = "pliron-mlir-export" +version = "0.1.0" +dependencies = [ + "pliron", + "pliron-derive", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" diff --git a/Cargo.toml b/Cargo.toml index 01ba4f4347..56595c3677 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,8 @@ members = [ "crates/cuda-host", "crates/cuda-macros", "crates/llvm-export", + "crates/pliron-mlir-export", + "crates/cuda-oxide-mlir-export", "crates/dialect-mir", "crates/dialect-iket", "crates/dialect-nvvm", @@ -67,6 +69,8 @@ pliron-llvm = { git = "https://github.com/pliron-org/pliron", rev = "edd41fe0615 # Internal crates cuda-target-spec = { path = "crates/cuda-target-spec" } llvm-export = { path = "crates/llvm-export" } +pliron-mlir-export = { path = "crates/pliron-mlir-export" } +cuda-oxide-mlir-export = { path = "crates/cuda-oxide-mlir-export" } cute-layout = { path = "cuteir/cute-layout" } dialect-cute = { path = "cuteir/dialect-cute" } dialect-mir = { path = "crates/dialect-mir" } @@ -113,6 +117,7 @@ combine = "4.6" 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/README.md b/README.md index ec171a87e0..e57a41978b 100644 --- a/README.md +++ b/README.md @@ -316,11 +316,13 @@ 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 | +| `pliron-mlir-export` | Typed deterministic Pliron-to-MLIR translation | +| `cuda-oxide-mlir-export` | CUTLASS 4.7 CuTe MLIR mapping packs and profile | | `dialect-nvvm` | pliron dialect modelling NVVM intrinsics | | `dialect-ptx` | pliron dialect modelling structured PTX | | `mir-transforms` | Optimization passes over the MIR dialect (loop unroll, ...) | | `nvvm-transforms` | Target-aware LLVM dialect legalization for NVVM | -| `cuda-oxide-codegen` | Rustc-independent PTX backend | +| `cuda-oxide-codegen` | Rustc-independent PTX and CUTLASS cubin backend | ### Build Tooling diff --git a/crates/cargo-oxide/src/commands/codegen_env.rs b/crates/cargo-oxide/src/commands/codegen_env.rs index 4bffb23d9c..1918c369fb 100644 --- a/crates/cargo-oxide/src/commands/codegen_env.rs +++ b/crates/cargo-oxide/src/commands/codegen_env.rs @@ -4,6 +4,7 @@ */ use crate::backend; +use crate::cutlass_toolchain; use std::path::Path; use std::process::Command; @@ -383,6 +384,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/crates/cargo-oxide/src/commands/fingerprint.rs b/crates/cargo-oxide/src/commands/fingerprint.rs index 03a9e4ce24..c37e7c3435 100644 --- a/crates/cargo-oxide/src/commands/fingerprint.rs +++ b/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/crates/cargo-oxide/src/commands/passthrough.rs b/crates/cargo-oxide/src/commands/passthrough.rs index b8c1a13988..4c3880c639 100644 --- a/crates/cargo-oxide/src/commands/passthrough.rs +++ b/crates/cargo-oxide/src/commands/passthrough.rs @@ -4,9 +4,11 @@ */ use std::collections::BTreeMap; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::process::Command; +use crate::cutlass_toolchain; + use super::*; /// Device debug-information policy requested on the command line. @@ -177,6 +179,79 @@ pub(super) fn inherited_process_env() -> BTreeMap> { .collect() } +pub(super) const DEVICE_BACKEND_ENV: &str = "CUDA_OXIDE_DEVICE_BACKEND"; +pub(super) const CUTLASS_DIGEST_FINGERPRINT_ENV: &str = + "CUDA_OXIDE_INTERNAL_CUTLASS_COMPILER_SHA256"; + +/// Include the selected CUTLASS compiler and its content identity in the +/// semantic environment seen by Cargo. +pub(super) fn inherited_process_env_with_cutlass_compiler_identity( + ctx: &Context, +) -> BTreeMap> { + let mut inherited = inherited_process_env(); + let inherited_compiler_path = + std::env::var_os(cutlass_toolchain::CUTLASS_COMPILER_ENV).map(PathBuf::from); + augment_env_with_cutlass_compiler_identity( + ctx, + &mut inherited, + inherited_compiler_path.as_deref(), + ); + inherited +} + +pub(super) fn augment_env_with_cutlass_compiler_identity( + ctx: &Context, + inherited: &mut BTreeMap>, + inherited_compiler_path: Option<&Path>, +) { + let backend = inherited + .get(DEVICE_BACKEND_ENV) + .map(Vec::as_slice) + .or_else(|| project_config_env(ctx, DEVICE_BACKEND_ENV).map(str::as_bytes)); + if !backend.is_some_and(selector_is_cutlass_mlir) { + return; + } + + let explicit_compiler = if inherited.contains_key(cutlass_toolchain::CUTLASS_COMPILER_ENV) { + inherited_compiler_path.map(Path::to_path_buf) + } else { + project_config_env(ctx, cutlass_toolchain::CUTLASS_COMPILER_ENV).map(PathBuf::from) + }; + if inherited.contains_key(cutlass_toolchain::CUTLASS_COMPILER_ENV) + || explicit_compiler.is_some() + { + let digest = match explicit_compiler.as_deref() { + Some(path) => cutlass_toolchain::sha256_file(path) + .unwrap_or_else(|error| format!("unavailable:{error}")), + None => "unavailable:non-unicode inherited compiler path".to_string(), + }; + inherited.insert( + CUTLASS_DIGEST_FINGERPRINT_ENV.to_string(), + digest.into_bytes(), + ); + return; + } + + if let Ok(Some(paths)) = cutlass_toolchain::resolve_official() { + inherited.insert( + cutlass_toolchain::CUTLASS_COMPILER_ENV.to_string(), + paths + .compiler_library + .as_os_str() + .as_encoded_bytes() + .to_vec(), + ); + inherited.insert( + CUTLASS_DIGEST_FINGERPRINT_ENV.to_string(), + paths.compiler_library_sha256.as_bytes().to_vec(), + ); + } +} + +pub(super) fn selector_is_cutlass_mlir(value: &[u8]) -> bool { + std::str::from_utf8(value).is_ok_and(|value| value.trim() == "cutlass-mlir") +} + fn cargo_passthrough_command( ctx: &Context, cargo_subcommand: CargoPassthroughSubcommand, diff --git a/crates/cargo-oxide/src/commands/tests.rs b/crates/cargo-oxide/src/commands/tests.rs index bd4e9dcb95..48986e6e9d 100644 --- a/crates/cargo-oxide/src/commands/tests.rs +++ b/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 @@ -68,6 +69,40 @@ fn is_sha256(value: &str) -> bool { .all(|byte| byte.is_ascii_digit() || matches!(byte, b'a'..=b'f')) } +#[test] +fn managed_cutlass_compiler_sets_only_the_library_env() { + let resolved = cutlass_toolchain::ResolvedCutlass { + install_dir: PathBuf::from("/managed/cutlass"), + compiler_library: PathBuf::from("/managed/cutlass/lib/libCutlassCompiler.so"), + compiler_library_sha256: "fixture-digest".into(), + }; + let mut cmd = Command::new("cargo"); + apply_resolved_cutlass_compiler(&mut cmd, Some(&resolved), false); + assert_eq!( + command_env(&cmd, cutlass_toolchain::CUTLASS_COMPILER_ENV).as_deref(), + Some("/managed/cutlass/lib/libCutlassCompiler.so") + ); +} + +#[test] +fn explicit_cutlass_compiler_outranks_managed_resolution() { + let resolved = cutlass_toolchain::ResolvedCutlass { + install_dir: PathBuf::from("/managed/cutlass"), + compiler_library: PathBuf::from("/managed/cutlass/lib/libCutlassCompiler.so"), + compiler_library_sha256: "fixture-digest".into(), + }; + let mut cmd = Command::new("cargo"); + cmd.env( + cutlass_toolchain::CUTLASS_COMPILER_ENV, + "/explicit/libCutlassCompiler.so", + ); + apply_resolved_cutlass_compiler(&mut cmd, Some(&resolved), true); + assert_eq!( + command_env(&cmd, cutlass_toolchain::CUTLASS_COMPILER_ENV).as_deref(), + Some("/explicit/libCutlassCompiler.so") + ); +} + /// `cargo_passthrough_command` with an empty ambient /// `CUDA_OXIDE_MATERIALIZE_CUBIN`. /// @@ -2941,6 +2976,100 @@ 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, + }; + 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, + ) + }; + 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, + }; + 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()); diff --git a/crates/cargo-oxide/src/cutlass_toolchain.rs b/crates/cargo-oxide/src/cutlass_toolchain.rs new file mode 100644 index 0000000000..67ad0349f5 --- /dev/null +++ b/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/crates/cargo-oxide/src/main.rs b/crates/cargo-oxide/src/main.rs index 16a6136aee..ecfff011d8 100644 --- a/crates/cargo-oxide/src/main.rs +++ b/crates/cargo-oxide/src/main.rs @@ -28,6 +28,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 +38,7 @@ mod artifact_identity; mod backend; mod backend_source; mod commands; +mod cutlass_toolchain; /// Top-level CLI structure parsed by clap. /// @@ -462,6 +464,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)] @@ -613,6 +635,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(), @@ -1092,6 +1118,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); + } + }, + }, } } @@ -1210,6 +1247,20 @@ mod tests { assert!(force); } + #[test] + fn toolchain_parser_accepts_pinned_cutlass_install() { + let cli = Cli::try_parse_from(["cargo-oxide", "toolchain", "install", "cutlass"]) + .expect("managed CUTLASS install command should parse"); + assert!(matches!( + cli.command, + Commands::Toolchain { + command: ToolchainCommand::Install { + tool: ManagedToolchain::Cutlass + } + } + )); + } + #[test] fn build_parser_preserves_nested_cargo_and_test_separators() { let args = strings(&[ diff --git a/crates/cuda-oxide-codegen/Cargo.toml b/crates/cuda-oxide-codegen/Cargo.toml index 8670166103..7bb7f268a6 100644 --- a/crates/cuda-oxide-codegen/Cargo.toml +++ b/crates/cuda-oxide-codegen/Cargo.toml @@ -19,11 +19,17 @@ 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 } llvm-export = { workspace = true } +pliron-mlir-export = { workspace = true } 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/crates/cuda-oxide-codegen/src/api.rs b/crates/cuda-oxide-codegen/src/api.rs index 2cce582b3b..6324dcf8c3 100644 --- a/crates/cuda-oxide-codegen/src/api.rs +++ b/crates/cuda-oxide-codegen/src/api.rs @@ -666,6 +666,7 @@ impl Compiler { let scratch = ScratchDirectory::new()?; let ll_path = scratch.path().join("module.ll"); let ptx_path = scratch.path().join("module.ptx"); + let cubin_path = scratch.path().join("module.cubin"); let backend_options = BackendOptions { iket: crate::options::IketInstrumentation::Auto, target_arch: Some(options.target.sm()), @@ -686,6 +687,8 @@ impl Compiler { OutputFiles { llvm_ir: &ll_path, ptx: &ptx_path, + cubin: &cubin_path, + mlir: None, stale_before_export: &[], }, ); @@ -1008,6 +1011,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/crates/cuda-oxide-codegen/src/cutlass_backend.rs b/crates/cuda-oxide-codegen/src/cutlass_backend.rs new file mode 100644 index 0000000000..fd5c2900ed --- /dev/null +++ b/crates/cuda-oxide-codegen/src/cutlass_backend.rs @@ -0,0 +1,1187 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Executable Backend-B boundary for the official CUTLASS compiler library. + +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}"; + +/// External input needed to execute the official CUTLASS Backend B. +#[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 the in-tree profile is accepted. + pub profile: String, +} + +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(), + } + } +} + +/// Device compiler selected after the shared MIR/CuTe preparation stages. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub enum DeviceBackend { + /// Ordinary CUDA Oxide MIR/NVVM/LLVM backend inherited from main. + #[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, + Scalar64, +} + +impl KernelParameterKind { + const fn size(self) -> u16 { + 8 + } + + const fn cuda_parameter_space(self) -> u32 { + match self { + Self::GlobalPointer64 => 5, + Self::Scalar64 => 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 config.profile != CUTLASS_FULL_CUTE_MLIR22_PROFILE { + return Err(backend_error(format!( + "unsupported MLIR profile {:?}; this compiler implements only {:?}", + config.profile, CUTLASS_FULL_CUTE_MLIR22_PROFILE + ))); + } + if expected_kernels.is_empty() { + return Err(backend_error( + "the CUTLASS backend requires at least one named kernel entry", + )); + } + parse_sm_target(target)?; + + 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, + )?; + 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, CUTLASS_PRECOMPILED_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, +) -> Result { + let pipeline_sha256 = format!("{:x}", Sha256::digest(CUTLASS_PRECOMPILED_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, + }); + 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) +} + +/// Recover the exact direct-launch ABI that the pinned export profile writes +/// on each kernel. Backend B currently flattens every Rust slice into a +/// `(pointer, i64)` pair and every remaining scalar into `i64`; accepting any +/// other type here would require a corresponding host-launch ABI decision. +fn expected_kernel_abis( + rendered_mlir: &str, + expected_kernels: &[String], +) -> Result { + const PREFIX: &str = "\"func.func\"() <{function_type = ("; + const RETURN_AND_NAME: &str = ") -> (), sym_name = \""; + + 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!("{RETURN_AND_NAME}{kernel}\""); + // The generic printer may append whitespace after the region opener; + // match the stable function/name portion instead of depending on it. + let mut signatures = rendered_mlir + .lines() + .filter_map(|line| { + let signature = line.trim().strip_prefix(PREFIX)?; + let end = signature.find(&name_marker)?; + signature[end + name_marker.len()..] + .trim_start() + .starts_with("}> ({") + .then_some(&signature[..end]) + }) + .collect::>(); + if signatures.len() != 1 { + return Err(format!( + "expected exactly one flattened func.func signature for kernel {kernel:?}, found {}", + signatures.len() + )); + } + let signature = signatures.pop().expect("length checked"); + let parameters = if signature.trim().is_empty() { + Vec::new() + } else { + signature + .split(',') + .map(|parameter| match parameter.trim() { + "!llvm.ptr" => Ok(KernelParameterKind::GlobalPointer64), + "i64" => Ok(KernelParameterKind::Scalar64), + other => Err(format!( + "kernel {kernel:?} has unsupported flattened parameter type {other:?}; only !llvm.ptr and i64 have a pinned direct-launch ABI" + )), + }) + .collect::, _>>()? + }; + abis.insert(kernel.clone(), parameters); + } + Ok(abis) +} + +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(()) +} + +/// Validate the post-ptxas direct-launch parameter layout recorded in CUDA's +/// per-kernel `.nv.info` section. These record tags and bit fields are part of +/// the fingerprint-pinned CUTLASS 4.7 / CUDA 13.x output contract; unknown or +/// malformed records are rejected instead of being guessed through. +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 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(()), + 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, + }; + + 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(); + 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(&((ordinal * 8) as u16).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(&((parameters.len() * 8) as u16).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}"); + } + + #[test] + fn flattened_signature_parser_rejects_unsupported_or_ambiguous_abis() { + let unsupported = + r#""func.func"() <{function_type = (!llvm.ptr, i32) -> (), 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("i32", "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/crates/cuda-oxide-codegen/src/cutlass_compiler.rs b/crates/cuda-oxide-codegen/src/cutlass_compiler.rs new file mode 100644 index 0000000000..13e2d0f5ff --- /dev/null +++ b/crates/cuda-oxide-codegen/src/cutlass_compiler.rs @@ -0,0 +1,896 @@ +/* + * 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. +//! +//! This module intentionally exposes one narrow operation: compile textual +//! `PreCompiledMlir` through a pinned CuTe pipeline to the C API's +//! `ObjectArtifact`. Despite containing device code, that artifact is a host +//! ELF object, not a standalone CUDA image. The +//! official API owns all compiler and artifact handles, so the wrappers below +//! release them through the matching C 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/crates/cuda-oxide-codegen/src/error.rs b/crates/cuda-oxide-codegen/src/error.rs index d47931b2dd..9febb1d74c 100644 --- a/crates/cuda-oxide-codegen/src/error.rs +++ b/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/crates/cuda-oxide-codegen/src/lib.rs b/crates/cuda-oxide-codegen/src/lib.rs index 75dd00e8fb..301732b88e 100644 --- a/crates/cuda-oxide-codegen/src/lib.rs +++ b/crates/cuda-oxide-codegen/src/lib.rs @@ -16,6 +16,8 @@ #![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/crates/cuda-oxide-codegen/src/pipeline.rs b/crates/cuda-oxide-codegen/src/pipeline.rs index 649ea57db7..be49b66b45 100644 --- a/crates/cuda-oxide-codegen/src/pipeline.rs +++ b/crates/cuda-oxide-codegen/src/pipeline.rs @@ -9,6 +9,7 @@ //! every destructive compiler stage after that boundary so the rustc frontend //! and the standalone frontend cannot silently diverge. +use crate::cutlass_backend::{CutlassBackendConfig, compile_to_cubin}; use crate::error::PipelineError; use crate::export::{ DebugExport, DeviceExternDecl, export_llvm_ir, function_local_static_placement_for_llc, @@ -30,11 +31,16 @@ use crate::ptx::{ }; use crate::target::detect_features_in_llvm_text; use crate::verify::verify_operation; +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, 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 pliron_mlir_export::render_module as render_mlir_module; +use std::ffi::OsString; +use std::io; use std::path::{Path, PathBuf}; /// Files used to publish one compiler result. @@ -45,6 +51,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], } @@ -99,16 +109,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 { @@ -121,6 +136,8 @@ impl<'a> ModulePipelineRequest<'a> { files, trace, generated_marker_policy: GeneratedMarkerPolicy::Required, + cutlass_backend, + expected_kernels, } } @@ -140,6 +157,8 @@ impl<'a> ModulePipelineRequest<'a> { files, trace: PipelineTrace::default(), generated_marker_policy: GeneratedMarkerPolicy::Optional, + cutlass_backend: None, + expected_kernels: &[], } } } @@ -152,6 +171,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. @@ -273,6 +294,83 @@ pub fn compile_translated_module( // materialization fails with its explicit-target requirement). materialize_iket(ctx, module, backend.target_arch.as_deref(), &backend.iket)?; + // Observe or compile the shared prepared semantic module before the + // ordinary inherited backend begins MIR/NVVM/LLVM lowering. Builds without + // an observation or CUTLASS selection do no translation-backend work. + let cutlass_rendered = if request.files.mlir.is_some() || request.cutlass_backend.is_some() { + Some(render_live_cutlass_mlir(ctx, module, backend)?) + } 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() { @@ -599,6 +697,207 @@ 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 render_live_cutlass_mlir( + ctx: &Context, + module: Ptr, + backend: &BackendOptions, +) -> Result { + let gpu_arch = backend + .target_arch + .as_deref() + .or(backend.device_arch_hint.as_deref()) + .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 = CutlassFullCuteMlir22::new(gpu_arch).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: &pliron_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'); + } + output.push_str(&render_mlir_module(module)); + Ok(output) +} + /// Backend decision made from the typed module before MIR lowering starts. #[derive(Clone, Copy, Debug, PartialEq, Eq)] struct PreLoweringBackendSelection { @@ -810,6 +1109,162 @@ fn remove_stale_files(paths: &[PathBuf]) -> Result<(), PipelineError> { mod tests { use super::*; + #[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_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(), + 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 = pliron_mlir_export::MlirModule { + root: pliron_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("//")) + ); + } + + #[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")); + } + fn typed_mir_test_module(ctx: &mut Context, callees: &[&str]) -> Ptr { use pliron::basic_block::BasicBlock; use pliron::builtin::attributes::{StringAttr, TypeAttr}; diff --git a/crates/cuda-oxide-mlir-export/Cargo.toml b/crates/cuda-oxide-mlir-export/Cargo.toml new file mode 100644 index 0000000000..e002335f5c --- /dev/null +++ b/crates/cuda-oxide-mlir-export/Cargo.toml @@ -0,0 +1,23 @@ +[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-export = { workspace = true } +dialect-mir = { workspace = true } +dialect-nvvm = { workspace = true } +dialect-cute = { workspace = true } +llvm-export = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } + +[dev-dependencies] +pliron-derive = { workspace = true } diff --git a/crates/cuda-oxide-mlir-export/README.md b/crates/cuda-oxide-mlir-export/README.md new file mode 100644 index 0000000000..b015cbddcb --- /dev/null +++ b/crates/cuda-oxide-mlir-export/README.md @@ -0,0 +1,157 @@ +# CUDA Oxide MLIR export + +This crate contains CUDA Oxide's target-specific mappings for +`pliron-mlir-export`. It is the CUTLASS translation backend. The generic +crate builds typed, deterministic MLIR text; this crate explains what each +CUDA Oxide operation means to one pinned MLIR consumer. + +```text +pliron-mlir-export this crate +------------------ ---------- +syntax tree + renderer + builtin / MIR / NVVM / CuTe mappings +registry + diagnostics pinned consumer profiles +``` + +The split is deliberate. The generic crate can translate any Pliron dialect. +This crate decides what CUDA Oxide operations mean in one exact MLIR consumer. + +The first profile targets the official CUTLASS 4.7 compiler library. CUDA +Oxide installs the exact NVIDIA release archive, verifies both the archive and +library digests, and loads that versioned library explicitly: + +```text +shared high-level Pliron module + │ + ▼ + shared CuTe safety checks + │ + ▼ + builtin + MIR + NVVM + CuTe recipes + │ + ▼ + generic MLIR renderer + │ + ▼ + official CUTLASS 4.7 compiler +``` + +The elementwise, NVFP4 GEMV, and block-scaled GEMM mapping packs translate the +live post-preparation module directly. The pinned library compiles that MLIR to +a validated cubin, which is embedded and launched through CUDA Oxide's ordinary +artifact loader. Set `CUDA_OXIDE_MLIR_OUTPUT=` to retain the exact textual +module for inspection; select `CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir` to compile +it. + +## The first mapping pack + +The first pack covers 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.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. + +The CUDA kernel marker also has an exact mapping: + +```text +gpu_kernel = "true" → cute.kernel +anything else → translation error +``` + +## The first CuTe slice + +The elementwise path keeps the useful layout explanation instead of reducing +everything to pointer arithmetic first: + +```text +make tensor + │ + ▼ +divide it into tiles + │ + ▼ +pick this thread's tile + │ + ├── full tile ──► cute.copy + │ + └── edge tile ──► checked scalar load/store +``` + +The profile recognizes the corresponding `cute.tensor_*` operations and maps +them to CUTLASS's `cute` and `cute_nvgpu` operations. Before translation, it +runs the shared backend-neutral whole-module verifier directly. The exporter +does not invoke native CuTe expansion or use native generated-intrinsic markers +as an intermediate representation. + +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. diff --git a/crates/cuda-oxide-mlir-export/src/cute.rs b/crates/cuda-oxide-mlir-export/src/cute.rs new file mode 100644 index 0000000000..e20b82c30d --- /dev/null +++ b/crates/cuda-oxide-mlir-export/src/cute.rs @@ -0,0 +1,1000 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! The first semantic CuTe mapping slice for the official CUTLASS `cute` +//! dialect. +//! +//! The split is deliberate: +//! +//! ```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 has first-class operations for the view transformations and bulk copy. +//! It does not have operations carrying the original logical length or the +//! absolute-index tail-store contract, so inventing similarly named target +//! operations would lose semantics rather than preserve them. + +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, +}; +use pliron_mlir_export::{ + DropAttribute, MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, + OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, TypeTranslation, +}; + +/// 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_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(()) +} + +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 element = scalar_spelling(®istry.translate_type(ctx, view.storage)?)?; + 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 + )) + } +} + +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 element = scalar_spelling(&session.translate_type(view.storage).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) + } +} + +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]) + } +} + +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]) + } +} + +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) + } +} + +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) + } +} + +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, + ) + } +} + +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, + ) + } +} + +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]) + } +} + +#[derive(Clone, Copy)] +struct TensorViewState { + data: Value, + len: Value, + tile_size: u64, + tile_index: Value, +} + +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}`" + )) +} + +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())?, + }) +} + +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)) +} + +#[derive(Clone, Copy)] +enum CopyDirection { + Load, + Store, +} + +#[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 element = scalar_spelling(&session.translate_type(view.storage).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) +} + +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(()) +} + +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()) +} + +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:?}" + )), + } +} + +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:?}" + )), + } +} + +fn scalar_spelling(ty: &MlirType) -> Result { + match ty { + MlirType::Integer(width) => Ok(format!("i{width}")), + MlirType::Float(pliron_mlir_export::MlirFloatType::F16) => Ok("f16".into()), + MlirType::Float(pliron_mlir_export::MlirFloatType::F32) => Ok("f32".into()), + other => Err(format!( + "the elementwise CUTLASS pack needs an i8, i16, f16, or f32 scalar, got {other:?}" + )), + } +} + +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_)) +} + +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}") + }) +} + +fn llvm_pointer_type(address_space: u32) -> Result { + if address_space == 0 { + MlirType::dialect("!llvm.ptr") + } else { + MlirType::dialect(format!("!llvm.ptr<{address_space}>")) + } +} + +fn cute_memref_type( + element: &str, + space: &str, + alignment: u64, + layout: &str, +) -> Result { + MlirType::dialect(format!( + "!cute.memref<{element}, {space}, align<{alignment}>, \"{layout}\">" + )) +} + +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(()) +} + +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) +} + +fn fresh(session: &mut TranslationSession<'_>, ty: MlirType) -> (MlirResult, MlirValueUse) { + let id = session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) +} + +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) +} + +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) +} + +fn binary( + name: &str, + result: MlirResult, + left: MlirValueUse, + right: MlirValueUse, + location: &MlirLocation, +) -> Result { + operation(name, vec![result], vec![left, right], location) +} + +fn unary( + name: &str, + result: MlirResult, + operand: MlirValueUse, + location: &MlirLocation, +) -> Result { + operation(name, vec![result], vec![operand], location) +} + +fn u64_bits(value: u64) -> i128 { + i64::from_ne_bytes(value.to_ne_bytes()) as i128 +} + +fn stringify(error: impl std::fmt::Display) -> String { + error.to_string() +} diff --git a/crates/cuda-oxide-mlir-export/src/cute_gemm.rs b/crates/cuda-oxide-mlir-export/src/cute_gemm.rs new file mode 100644 index 0000000000..f7a090f5e4 --- /dev/null +++ b/crates/cuda-oxide-mlir-export/src/cute_gemm.rs @@ -0,0 +1,3190 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SM120 block-scaled GEMM mappings for the official CUTLASS 4.7 compiler. +//! +//! The source dialect keeps scheduler, pipeline, TMA, MMA, and epilogue +//! contracts visible. This pack erases only compiler-only handles and lowers +//! their runtime leaves to operations accepted by the public PRE_COMPILED +//! pipeline in `libCutlassCompiler.so`. + +use dialect_cute::{ + attributes::{ + CuteCountedCtaBarrierAttr, CuteEpilogueHalfAttr, CuteEpilogueSyncPhaseAttr, + CuteMmaCarrierKindAttr, CutePipelineRoleAttr, CutePipelineStateAttr, + CuteTensorAddressSpaceAttr, 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::{MirConstructTupleOp, MirExternSharedOp, MirInsertFieldOp}, + types::{MirArrayType, MirStructType, MirTupleType}, +}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, + r#type::{TypeHandle, Typed}, +}; +use pliron_mlir_export::{ + DropAttribute, MlirAttribute, MlirBlock, MlirLocation, MlirOperation, MlirRegion, MlirResult, + MlirType, MlirValueUse, OperationInput, OperationTranslation, TranslationError, + TranslationRegistry, TranslationSession, TypeTranslation, +}; + +const PIPELINE_WAIT_TICKS: i128 = 10_000_000; +const LDMATRIX_BYTE_SWIZZLE_MASK: u64 = 0x180; +const LDMATRIX_BYTE_SWIZZLE_SHIFT: u64 = 3; +const COUNTED_CTA_SYNC_OP: &str = "nvvm.barrier.cta.sync"; + +/// 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::(ConstructAggregateTranslation)?; + 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) +} + +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)) + } +} + +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() + } +} + +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(), + } + } +} + +struct StaticContractTypeTranslation; + +impl TypeTranslation for StaticContractTypeTranslation { + fn translate( + &self, + _ctx: &Context, + _source: TypeHandle, + _registry: &TranslationRegistry, + ) -> Result { + Ok(MlirType::Tuple(vec![])) + } +} + +struct ConstructAggregateTranslation; + +impl OperationTranslation for ConstructAggregateTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + MirConstructTupleOp::new(source) + .verify(ctx) + .map_err(stringify)?; + expect_shape("mir.construct_tuple", &input, 1, input.operands.len())?; + build_aggregate( + session, + input.results[0].clone(), + input.operands, + &input.location, + ) + } +} + +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; + let mut target = operation( + "llvm.insertvalue", + input.results, + input.operands, + &input.location, + )?; + target + .properties + .insert("position".into(), MlirAttribute::DenseI64Array(vec![index])); + Ok(vec![target]) + } +} + +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]) + } +} + +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) + } +} + +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]) + } +} + +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) + } +} + +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) + } +} + +struct IdentityTranslation(&'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, + )?]) + } +} + +struct IdentityFirstOperandTranslation(&'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. + +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) + } +} + +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, + )?, + ]) + } +} + +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, + )?]) + } +} + +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) + } +} + +struct PipelineWaitTranslation { + empty_ring: bool, + 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) + } +} + +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) + } +} + +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) + } +} + +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) + } +} + +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) + } +} + +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) + } +} + +struct TmaStoreWaitTranslation { + 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]) + } +} + +struct NoOperandRenameTranslation(&'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)?]) + } +} + +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) + } +} + +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) + } +} + +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(()) +} + +struct StaticMultiIdentityTranslation { + name: &'static str, + 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() + } +} + +#[derive(Clone)] +struct CarrierLayout { + target: MlirType, + kind: CarrierKind, +} + +#[derive(Clone)] +enum CarrierKind { + Leaf, + Aggregate(Vec<(i64, CarrierLayout)>), +} + +impl CarrierLayout { + 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), + }) + } + + fn leaf_count(&self) -> usize { + match &self.kind { + CarrierKind::Leaf => 1, + CarrierKind::Aggregate(children) => { + children.iter().map(|(_, child)| child.leaf_count()).sum() + } + } + } +} + +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(()) + } + } +} + +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(()) +} + +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) +} + +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) + } +} + +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) + } +} + +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) + } +} + +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) + } +} + +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()) +} + +struct TiledGemmTranslation; + +impl OperationTranslation for TiledGemmTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CuteTiledGemmOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + 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(), + &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) + } +} + +#[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, + 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 = cast_int( + session, + operations, + scale_a, + MlirType::Integer(16), + "arith.trunci", + location, + )?; + let scale_b = 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(pliron_mlir_export::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(pliron_mlir_export::MlirFloatType::Other( + "f4E2M1FN".into(), + ))), + ); + operation_.properties.insert( + "b_type".into(), + MlirAttribute::Type(MlirType::Float(pliron_mlir_export::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(pliron_mlir_export::MlirFloatType::Other( + "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: 32, + ty: MlirType::Integer(32), + }, + ); + operations.push(operation_); + uses.try_into() + .map_err(|_| "SM120 MXF4 MMA did not produce four accumulators".to_owned()) +} + +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)?; + + 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) + } +} + +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(pliron_mlir_export::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) +} + +fn f16x2_intrinsic_operands(low: MlirValueUse, high: MlirValueUse) -> [MlirValueUse; 2] { + [high, low] +} + +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_) +} + +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_) +} + +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()) +} + +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.inner().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)) +} + +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_) +} + +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, + ) +} + +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, + ) +} + +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])) +} + +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_) +} + +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), + ); +} + +fn build_aggregate( + session: &mut TranslationSession<'_>, + final_result: MlirResult, + fields: Vec, + location: &MlirLocation, +) -> Result, String> { + if fields.is_empty() { + return Ok(vec![operation( + "llvm.mlir.poison", + vec![final_result], + vec![], + location, + )?]); + } + let (poison_result, mut aggregate) = fresh(session, final_result.ty.clone()); + let mut operations = vec![operation( + "llvm.mlir.poison", + vec![poison_result], + vec![], + location, + )?]; + 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(), + }; + let mut insert = operation( + "llvm.insertvalue", + vec![result], + vec![aggregate, field], + location, + )?; + insert.properties.insert( + "position".into(), + MlirAttribute::DenseI64Array(vec![position as i64]), + ); + operations.push(insert); + aggregate = next; + } + Ok(operations) +} + +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_) +} + +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_) +} + +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_)) +} + +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_) +} + +fn binary( + name: &str, + result: MlirResult, + left: MlirValueUse, + right: MlirValueUse, + location: &MlirLocation, +) -> Result { + operation(name, vec![result], vec![left, right], location) +} + +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) +} + +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(()) +} + +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) +} + +fn mbarrier_init( + barrier: MlirValueUse, + arrival_count: MlirValueUse, + location: &MlirLocation, +) -> Result { + operation( + "nvvm.mbarrier.init", + vec![], + vec![barrier, arrival_count], + location, + ) +} + +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, + ) +} + +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) +} + +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) +} + +fn epilogue_gep_element_type() -> MlirType { + MlirType::Float(pliron_mlir_export::MlirFloatType::F16) +} + +#[cfg(test)] +const fn ldmatrix_swizzled_byte_offset(plain: u64) -> u64 { + plain ^ ((plain & LDMATRIX_BYTE_SWIZZLE_MASK) >> LDMATRIX_BYTE_SWIZZLE_SHIFT) +} + +fn fresh(session: &mut TranslationSession<'_>, ty: MlirType) -> (MlirResult, MlirValueUse) { + let id = session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) +} + +fn generic_pointer_type() -> Result { + MlirType::dialect("!llvm.ptr") +} + +fn shared_pointer_type() -> Result { + MlirType::dialect("!llvm.ptr<3>") +} + +fn u64_bits(value: u64) -> i128 { + i64::from_ne_bytes(value.to_ne_bytes()) as i128 +} + +fn stringify(error: impl std::fmt::Display) -> String { + error.to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + use pliron_mlir_export::{MlirFloatType, MlirValueId}; + + fn value(id: u64, ty: MlirType) -> MlirValueUse { + MlirValueUse { + id: MlirValueId(id), + ty, + } + } + + #[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); + } + + #[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); + } + + #[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) + ); + } + + #[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/crates/cuda-oxide-mlir-export/src/cute_gemv.rs b/crates/cuda-oxide-mlir-export/src/cute_gemv.rs new file mode 100644 index 0000000000..5eadd84bda --- /dev/null +++ b/crates/cuda-oxide-mlir-export/src/cute_gemv.rs @@ -0,0 +1,1243 @@ +/* + * 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. +//! +//! The source views are compiler-only values, but the target carriers are +//! ordinary LLVM aggregates. Keeping the runtime pointer and shape fields in +//! SSA gives each selection operation a real target dataflow edge: +//! +//! ```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 +//! ``` +//! +//! This translation consumes the shared semantic operations directly and +//! manufactures only target operations owned by CUTLASS. + +use crate::cute::{TensorViewTypeTranslation, register_cute_elementwise_mappings}; + +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}, +}; +use pliron_mlir_export::{ + MlirAttribute, MlirFloatType, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, + OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, TypeTranslation, +}; + +const VALUE_BYTES_PER_TILE: u64 = 32; +const VALUE_BYTES_PER_LOAD: u64 = 16; +const VALUE_LANES_PER_LOAD: u64 = 32; +const SCALE_BYTES_PER_TILE: u64 = 4; +const SCALE_ATOM_ROWS: u64 = 128; +const SCALE_ATOM_K: u64 = 64; +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) +} + +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(()) +} + +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) + } +} + +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() + } +} + +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() + } +} + +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, + ) + } +} + +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) + } +} + +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]) + } +} + +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, + )?]) + } +} + +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()); + } + + emit_scaled_k64_load( + session, + input.operands[0].clone(), + input.results[0].clone(), + value_alignment, + scale_alignment, + &input.location, + ) + } +} + +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, + ) + } +} + +fn emit_scaled_k64_load( + session: &mut TranslationSession<'_>, + scaled: MlirValueUse, + final_result: MlirResult, + value_alignment: u64, + scale_alignment: u64, + 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, + NarrowLoad::UE8M0, + scale_alignment, + location, + )?; + + operations.extend(build_aggregate( + session, + final_result, + vec![value_lo, value_hi, scale_values], + location, + )?); + Ok(operations) +} + +#[derive(Clone, Copy)] +enum NarrowLoad { + E2M1, + UE8M0, +} + +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"), + }; + // 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) +} + +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) +} + +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, + ) +} + +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_) +} + +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_) +} + +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_) +} + +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) +} + +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) +} + +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) +} + +fn tensor_carrier_type() -> Result { + MlirType::dialect("!llvm.struct<(!llvm.ptr, i64, i64, i64)>") +} + +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)>", + ) +} + +fn fragment_carrier_type() -> Result { + MlirType::dialect("!llvm.struct<(vector<32xf32>, vector<32xf32>, vector<4xf32>)>") +} + +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())), + ), + } +} + +fn narrow_float(spelling: &str) -> MlirType { + match spelling { + "f16" => MlirType::Float(MlirFloatType::F16), + "bf16" => MlirType::Float(MlirFloatType::BF16), + "f4E2M1FN" | "f8E8M0FNU" => MlirType::Float(MlirFloatType::Other(spelling.into())), + _ => unreachable!("fixed GEMV conversion intermediate"), + } +} + +fn vector_type(size: u64, element: MlirType) -> MlirType { + MlirType::Vector { + shape: vec![size], + element: Box::new(element), + } +} + +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()) +} + +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()) +} + +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()) +} + +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, CuteTensorLayoutAttr::BlockScaleKMajor(16)) => Ok(()), + other => Err(format!( + "GEMV tensor needs E2M1/KMajor or UE8M0/BlockScaleKMajor<16>, got {other:?}" + )), + } +} + +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(()) +} + +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) +} + +fn fresh(session: &mut TranslationSession<'_>, ty: MlirType) -> (MlirResult, MlirValueUse) { + let id = session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) +} + +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) +} + +fn binary( + name: &str, + result: MlirResult, + left: MlirValueUse, + right: MlirValueUse, + location: &MlirLocation, +) -> Result { + operation(name, vec![result], vec![left, right], location) +} + +fn llvm_pointer_type(address_space: u32) -> Result { + if address_space == 0 { + MlirType::dialect("!llvm.ptr") + } else { + MlirType::dialect(format!("!llvm.ptr<{address_space}>")) + } +} + +fn u64_bits(value: u64) -> i128 { + i64::from_ne_bytes(value.to_ne_bytes()) as i128 +} + +fn stringify(error: impl std::fmt::Display) -> String { + error.to_string() +} diff --git a/crates/cuda-oxide-mlir-export/src/lib.rs b/crates/cuda-oxide-mlir-export/src/lib.rs new file mode 100644 index 0000000000..fee9654554 --- /dev/null +++ b/crates/cuda-oxide-mlir-export/src/lib.rs @@ -0,0 +1,27 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! CUDA Oxide mappings for textual MLIR export. +//! +//! The generic exporter answers "how do we build typed, deterministic MLIR +//! text?" This crate answers "what does this CUDA Oxide operation become for +//! this exact consumer?" + +mod cute; +mod cute_gemm; +mod cute_gemv; +mod mir_core; +mod mir_memory; +mod nvvm_sregs; +mod packs; +mod profile; + +pub use cute::register_cute_elementwise_pack; +pub(crate) use cute_gemm::register_cute_gemm_pack; +pub(crate) use cute_gemv::register_cute_cutlass_profile_pack; +pub use mir_core::register_mir_core_pack; +pub use nvvm_sregs::register_nvvm_sreg_pack; +pub use packs::register_builtin_pack; +pub use profile::{CutlassFullCuteMlir22, ExportManifest, MlirConsumerProfile, ProfileError}; diff --git a/crates/cuda-oxide-mlir-export/src/mir_core.rs b/crates/cuda-oxide-mlir-export/src/mir_core.rs new file mode 100644 index 0000000000..ec5ec8ea3a --- /dev/null +++ b/crates/cuda-oxide-mlir-export/src/mir_core.rs @@ -0,0 +1,1273 @@ +/* + * 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 dialect_mir::{ + attributes::{MirCastKindAttr, MirFP16Attr, MirPointerKindAuthorityAttr}, + 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::OperandSegmentInterface, + 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, +}; +use pliron_mlir_export::{ + AttributeTranslation, DropAttribute, FixedType, MlirAttribute, MlirBlockArgument, + MlirFloatType, MlirOperation, MlirResult, MlirType, MlirValueUse, OperationInput, + OperationTranslation, TranslationError, TranslationRegistry, TranslationSession, +}; + +/// 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)?; + 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)?; + + 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::(IntegerBinaryTranslation("arith.shli"))?; + registry.register_operation::(ShiftRightTranslation)?; + 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)?; + Ok(()) +} + +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), + })) + } +} + +struct RenameOperation(&'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)?]) + } +} + +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, "sym_name", "sym_name")?; + move_property(&mut target, "mir_func_type", "function_type")?; + lower_alwaysinline_contract(&mut target)?; + if let Some(marker) = target.attributes.remove("gpu_kernel") { + match marker { + MlirAttribute::String(value) if value == "true" => { + target + .attributes + .insert("cute.kernel".into(), MlirAttribute::Unit); + target + .attributes + .insert("gpu.kernel".into(), MlirAttribute::Unit); + move_exact_block_contract(&mut target)?; + flatten_kernel_slice_abi(ctx, source, &mut target, session)?; + } + other => { + return Err(format!( + "gpu_kernel must be the string `true`, got {other:?}" + )); + } + } + } + Ok(vec![target]) + } +} + +/// 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 is safe and useful for the pointer-free arithmetic +/// helpers in this profile. Pointer-bearing helpers are deliberately left to +/// CUTLASS 4.7's normal inliner because forcing them across address-space +/// boundaries is not part of this contract. +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) = function_type else { + return Err("alwaysinline function_type is not a type".into()); + }; + if mlir_type_contains_pointer(function_type) { + return Ok(()); + } + 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:?}" + )), + } +} + +fn mlir_type_contains_pointer(ty: &MlirType) -> bool { + match ty { + MlirType::Function { inputs, results } => { + inputs.iter().chain(results).any(mlir_type_contains_pointer) + } + MlirType::Vector { element, .. } => mlir_type_contains_pointer(element), + MlirType::Tuple(fields) => fields.iter().any(mlir_type_contains_pointer), + MlirType::Dialect(spelling) => { + spelling.contains("!llvm.ptr") || spelling.contains("!cute.ptr") + } + MlirType::Index | MlirType::Integer(_) | MlirType::Float(_) => false, + } +} + +/// 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()), + } +} + +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)) +} + +/// Match CUDA Oxide's existing kernel launch ABI: a Rust slice is passed as +/// two driver arguments, pointer then length. Rebuild the source aggregate at +/// entry so the already-translated body does not need any rewriting. +fn flatten_kernel_slice_abi( + ctx: &Context, + source: Ptr, + target: &mut MlirOperation, + session: &mut TranslationSession<'_>, +) -> 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(); + + for ((source_type, target_type), original_argument) in source_inputs + .into_iter() + .zip(target_inputs.iter()) + .zip(entry.arguments.iter()) + { + 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()); + 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()]); + 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; + Ok(()) +} + +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]) + } +} + +enum ConstantTranslation { + Integer, + 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]) + } +} + +enum NumericBinaryTranslation { + Add, + Sub, + Mul, + Div, + 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)?]) + } +} + +struct IntegerBinaryTranslation(&'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)?]) + } +} + +struct ShiftRightTranslation; + +impl OperationTranslation for ShiftRightTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let integer = source_integer(ctx, source, 0)?; + let name = if integer.signedness() == Signedness::Signed { + "arith.shrsi" + } else { + "arith.shrui" + }; + Ok(vec![renamed(name, input)?]) + } +} + +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(pliron_mlir_export::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, + pliron_mlir_export::MlirValueUse { + id: zero_id, + ty: result_type, + }, + ); + Ok(vec![zero, neg]) + } + } + } +} + +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(pliron_mlir_export::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(pliron_mlir_export::MlirValueUse { + id: ones_id, + ty: result_type, + }); + Ok(vec![ones, not]) + } +} + +enum ComparisonTranslation { + Lt, + Le, + Gt, + Ge, + Eq, + 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]) + } +} + +struct CastTranslation; + +impl OperationTranslation for CastTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + 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(); + 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 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)?]) + } +} + +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()) +} + +#[derive(Clone, Copy)] +enum Scalar { + Integer { width: u32, signed: bool }, + Float { width: u32 }, +} + +impl Scalar { + fn width(self) -> u32 { + match self { + Self::Integer { width, .. } | Self::Float { width } => width, + } + } +} + +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) +} + +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) +} + +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 }) +} + +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()) +} + +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) +} + +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 { + use std::num::NonZero; + + 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, + }; + + fn context() -> Context { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + ctx + } + + 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) + } + + 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 + } + + 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(), + ) + } + + fn integer(ctx: &mut Context, width: u32, signedness: Signedness) -> TypedHandle { + IntegerType::get(ctx, width, signedness) + } + + #[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) -> () + }) : () -> () +}) : () -> () +"#; + assert_eq!(text, expected); + } + + #[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}" + ); + } + + #[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}"); + } + + #[test] + fn alwaysinline_pointer_helper_is_left_to_cutlass_inlining() { + 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("passthrough"), "{text}"); + assert!(!text.contains(r#"alwaysinline = "true""#), "{text}"); + } + + #[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) -> () + }) : () -> () +}) : () -> () +"#; + assert_eq!(text, expected); + } +} diff --git a/crates/cuda-oxide-mlir-export/src/mir_memory.rs b/crates/cuda-oxide-mlir-export/src/mir_memory.rs new file mode 100644 index 0000000000..57add5c211 --- /dev/null +++ b/crates/cuda-oxide-mlir-export/src/mir_memory.rs @@ -0,0 +1,1308 @@ +/* + * 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 dialect_mir::{ + attributes::FieldIndexAttr, + ops::{ + MirAllocaOp, MirArrayElementAddrOp, MirAssertOp, MirCallOp, MirConstructStructOp, + MirExtractArrayElementOp, MirExtractFieldOp, MirFieldAddrOp, MirLoadOp, MirPtrOffsetOp, + MirStorageDeadOp, MirStorageLiveOp, MirStoreOp, + }, + types::{ + MirArrayType, MirDisjointSliceType, MirFP16Type, MirPtrType, MirSliceType, MirStructType, + MirTupleType, + }, +}; +use llvm_export::ops::LocalMemoryProvenanceAttr; +use pliron::{ + builtin::types::{FP16Type, FP32Type, FP64Type, IntegerType, UnitType}, + context::{Context, Ptr}, + operation::Operation, + r#type::{TypeHandle, Typed}, +}; +use pliron_mlir_export::{ + DropAttribute, MlirAttribute, MlirFloatType, MlirOperation, MlirResult, MlirType, MlirValueUse, + OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, TypeTranslation, +}; + +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)?; + + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + + registry.register_operation::(AllocaTranslation)?; + registry.register_operation::(LoadTranslation)?; + registry.register_operation::(StoreTranslation)?; + registry.register_operation::(PointerOffsetTranslation)?; + registry.register_operation::(ExtractFieldTranslation)?; + registry.register_operation::(ConstructStructTranslation)?; + 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(()) +} + +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()) + } +} + +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)]) + } +} + +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) + } +} + +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())?, + ) + } +} + +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) + } +} + +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) + } +} + +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]) + } +} + +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]) + } +} + +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]) + } +} + +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]) + } +} + +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.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]) + } +} + +struct ConstructStructTranslation; + +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) + } +} + +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))?, + }; + 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]) + } +} + +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]) + } +} + +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)?; + 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]) + } +} + +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:?}")); + } + let mut target = renamed("func.call", input)?; + target + .properties + .insert("callee".into(), MlirAttribute::SymbolRef(callee)); + Ok(vec![target]) + } +} + +struct AssertTranslation; + +impl OperationTranslation for AssertTranslation { + fn translate( + &self, + _ctx: &Context, + _source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + ensure_empty_attributes("mir.assert", &input)?; + if input.operands.is_empty() || input.successors.len() != 1 { + return Err("mir.assert expected a condition and one success block".into()); + } + let mut assertion = MlirOperation::new("cf.assert")?; + assertion.operands.push(input.operands[0].clone()); + assertion.properties.insert( + "msg".into(), + MlirAttribute::String("CUDA Oxide bounds assertion failed".into()), + ); + assertion.location = input.location.clone(); + + let mut branch = MlirOperation::new("cf.br")?; + branch.operands.extend_from_slice(&input.operands[1..]); + branch.successors = input.successors; + branch.location = input.location; + Ok(vec![assertion, branch]) + } +} + +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![]) + } +} + +#[derive(Debug)] +pub(crate) struct AggregatePlan { + pub(crate) ty: MlirType, + pub(crate) decl_to_slot: Vec>, + pub(crate) memory_order: Vec, +} + +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, + }) +} + +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(()) +} + +fn source_is_zero_sized(ctx: &Context, source: TypeHandle) -> bool { + let source_ref = source.deref(ctx); + 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 +} + +fn source_stored_size(ctx: &Context, source: TypeHandle) -> Option { + let source_ref = source.deref(ctx); + 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 +} + +fn source_natural_alignment(ctx: &Context, source: TypeHandle) -> Option { + let source_ref = source.deref(ctx); + 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 +} + +fn source_abi_align(ctx: &Context, source: TypeHandle) -> Option { + let source_ref = source.deref(ctx); + 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) +} + +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() +} + +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()) +} + +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()) +} + +fn llvm_pointer_type(address_space: u32) -> Result { + if address_space == 0 { + MlirType::dialect("!llvm.ptr") + } else { + MlirType::dialect(format!("!llvm.ptr<{address_space}>")) + } +} + +fn llvm_array_type(size: u64, element: &MlirType) -> Result { + MlirType::dialect(format!("!llvm.array<{size} x {}>", type_spelling(element))) +} + +fn llvm_struct_type(fields: &[MlirType]) -> Result { + let fields = fields + .iter() + .map(type_spelling) + .collect::>() + .join(", "); + MlirType::dialect(format!("!llvm.struct<({fields})>")) +} + +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(), + } +} + +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())), + } +} + +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::>() + )) + } +} + +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) +} + +fn memory_ordering(operation: &mut MlirOperation) { + operation.properties.insert( + "ordering".into(), + MlirAttribute::Integer { + value: 0, + ty: MlirType::Integer(64), + }, + ); +} + +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), + }, + ); + } +} + +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 { + use crate::{ + CutlassFullCuteMlir22, MlirConsumerProfile, + profile::render_mapping_module_without_cutlass_envelope, + }; + use dialect_mir::{ + attributes::FieldIndexAttr, + ops::{ + MirAddOp, MirAllocaOp, MirArrayElementAddrOp, 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, + }; + + 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 + } + + fn integer(ctx: &mut Context, width: u32) -> TypedHandle { + IntegerType::get(ctx, width, Signedness::Unsigned) + } + + #[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!(text, expected); + } +} diff --git a/crates/cuda-oxide-mlir-export/src/nvvm_sregs.rs b/crates/cuda-oxide-mlir-export/src/nvvm_sregs.rs new file mode 100644 index 0000000000..ec303a5849 --- /dev/null +++ b/crates/cuda-oxide-mlir-export/src/nvvm_sregs.rs @@ -0,0 +1,256 @@ +/* + * 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 dialect_nvvm::ops::{ + ReadPtxSregCtaidXOp, ReadPtxSregCtaidYOp, ReadPtxSregCtaidZOp, ReadPtxSregNctaidXOp, + ReadPtxSregNctaidYOp, ReadPtxSregNctaidZOp, ReadPtxSregNtidXOp, ReadPtxSregNtidYOp, + ReadPtxSregNtidZOp, ReadPtxSregTidXOp, ReadPtxSregTidYOp, ReadPtxSregTidZOp, +}; +use pliron::{context::Context, context::Ptr, operation::Operation}; +use pliron_mlir_export::{ + MlirOperation, OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, +}; + +/// 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(()) +} + +struct SregTranslation { + target: &'static str, + 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 != pliron_mlir_export::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 { + 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, + }; + + 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 + } + + #[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"() : () -> () + }) : () -> () +}) : () -> () +"#; + assert_eq!(text, 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/crates/cuda-oxide-mlir-export/src/packs.rs b/crates/cuda-oxide-mlir-export/src/packs.rs new file mode 100644 index 0000000000..2b68b9e866 --- /dev/null +++ b/crates/cuda-oxide-mlir-export/src/packs.rs @@ -0,0 +1,300 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use pliron::{ + attribute::AttrObj, + builtin::{ + attr_interfaces::TypedAttrInterface, + attributes::{ + BoolAttr, FPDoubleAttr, FPHalfAttr, FPSingleAttr, IdentifierAttr, IntegerAttr, + OperandSegmentSizesAttr, StringAttr, TypeAttr, UnitAttr, + }, + ops::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, +}; +use pliron_mlir_export::{ + AttributeTranslation, DropAttribute, FixedType, MlirAttribute, MlirFloatType, MlirOperation, + MlirType, OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, TypeTranslation, +}; + +/// 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_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(()) +} + +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("sym_name") { + target.properties.insert("sym_name".into(), name); + } + Ok(vec![target]) + } +} + +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())) + } +} + +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 }) + } +} + +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![])) + } +} + +fn translate_nested_type( + ctx: &Context, + registry: &TranslationRegistry, + source: TypeHandle, +) -> Result { + registry.translate_type(ctx, source) +} + +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()))) + } +} + +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()))) + } +} + +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()))) + } +} + +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)?, + })) + } +} + +enum FloatAttributeTranslation { + F16, + F32, + 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 })) + } +} + +struct UnitAttributeTranslation; + +impl AttributeTranslation for UnitAttributeTranslation { + fn translate( + &self, + _ctx: &Context, + _source: &AttrObj, + _registry: &TranslationRegistry, + ) -> Result, String> { + Ok(Some(MlirAttribute::Unit)) + } +} + +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/crates/cuda-oxide-mlir-export/src/profile.rs b/crates/cuda-oxide-mlir-export/src/profile.rs new file mode 100644 index 0000000000..f8faf81150 --- /dev/null +++ b/crates/cuda-oxide-mlir-export/src/profile.rs @@ -0,0 +1,273 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use pliron::{builtin::ops::ModuleOp, context::Context, op::Op}; +use pliron_mlir_export::{ + MlirAttribute, MlirBlock, MlirBlockId, MlirModule, MlirOperation, MlirRegion, + TranslationConfig, TranslationError, TranslationRegistry, + translate_module as translate_pliron_module, +}; +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +use crate::{ + register_builtin_pack, register_cute_cutlass_profile_pack, register_cute_gemm_pack, + register_mir_core_pack, register_nvvm_sreg_pack, +}; + +/// Everything that makes one textual MLIR contract reproducible. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ExportManifest { + pub profile: String, + pub exporter_version: String, + pub pliron_revision: String, + pub consumer_revision: String, + pub llvm_major: u32, + pub gpu_arch: String, + pub pointer_model: String, + pub index_bits: u32, +} + +impl ExportManifest { + pub fn to_pretty_json(&self) -> Result { + serde_json::to_string_pretty(self) + } +} + +#[derive(Debug, Error)] +pub enum ProfileError { + #[error(transparent)] + Registry(#[from] TranslationError), + #[error("CuTe source validation failed before MLIR translation: {0}")] + CuteSemantics(String), + #[error("could not build CUTLASS's executable GPU module envelope: {0}")] + ExecutableEnvelope(String), + #[error("unsupported GPU architecture `{0}` for this MLIR profile")] + UnsupportedArchitecture(String), +} + +/// A pinned target contract plus the mapping packs that implement it. +pub trait MlirConsumerProfile { + fn name(&self) -> &'static str; + fn manifest(&self) -> ExportManifest; + 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>; + + 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(), + )?) + } +} + +/// Full CuTe dialect accepted by the pinned official CUTLASS 4.7 compiler. +/// +/// This is intentionally versioned. A different public CUTLASS release is a +/// different profile until its syntax, lowering, and runtime gates pass. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CutlassFullCuteMlir22 { + gpu_arch: String, +} + +impl CutlassFullCuteMlir22 { + pub const CONSUMER_REVISION: &'static str = "cutlass-compiler-v4.7.0"; + pub const PLIRON_REVISION: &'static str = "8447aa426e2f090dc8b7b9383c5035db7bc70b62"; + + 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(), + 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_cute_cutlass_profile_pack(&mut registry)?; + register_cute_gemm_pack(&mut registry)?; + registry.seal(); + Ok(registry) + } + + fn verify_source_module(&self, ctx: &Context, module: &ModuleOp) -> Result<(), ProfileError> { + // Every backend enters through the same immutable whole-module CuTe + // graph/provenance checks before selecting its continuation. + dialect_cute::verify::verify_cute_semantics(ctx, module.get_operation()) + .map_err(|error| ProfileError::CuteSemantics(error.to_string())) + } + + 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) + } +} + +/// 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. +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 device_operations = std::mem::take(&mut root_block.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(()) +} + +fn next_block_id(operation: &MlirOperation) -> Result { + 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())); + pliron_mlir_export::render_module(&unwrapped) +} + +#[cfg(test)] +mod tests { + use super::*; + use pliron::identifier::Identifier; + + #[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("\"llvm_major\": 22")); + assert!(profile.build_registry().unwrap().is_sealed()); + } + + #[test] + fn profile_rejects_an_unproved_architecture() { + assert!(matches!( + CutlassFullCuteMlir22::new("sm_90a"), + Err(ProfileError::UnsupportedArchitecture(arch)) if arch == "sm_90a" + )); + } + + #[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/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs b/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs new file mode 100644 index 0000000000..8494a7d568 --- /dev/null +++ b/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs @@ -0,0 +1,439 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use std::num::NonZero; + +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, +}; +use pliron_mlir_export::render_module; + +fn append(block: Ptr, operation: O, ctx: &Context) -> Ptr { + let operation = operation.get_operation(); + operation.insert_at_back(block, ctx); + operation +} + +#[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); + + 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}"); +} + +#[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 [("reqntid_x", 256), ("reqntid_y", 1), ("reqntid_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(&pliron_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(&pliron_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(&pliron_mlir_export::MlirAttribute::Unit) + ); + assert_eq!( + kernel.attributes.get("gpu.kernel"), + Some(&pliron_mlir_export::MlirAttribute::Unit) + ); + assert_eq!( + kernel.attributes.get("nvvm.reqntid"), + Some(&pliron_mlir_export::MlirAttribute::DenseI32Array(vec![ + 256, 1, 1 + ])) + ); + let pliron_mlir_export::MlirAttribute::Type(pliron_mlir_export::MlirType::Function { + inputs, + .. + }) = kernel.properties.get("function_type").unwrap() + else { + panic!("kernel function_type is not a function type") + }; + let pointer = pliron_mlir_export::MlirType::dialect("!llvm.ptr").unwrap(); + let length = pliron_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); + assert_eq!(text.matches("\"gpu.module\"").count(), 1, "{text}"); + assert!(text.contains("gpu.container_module = unit"), "{text}"); + assert!( + text.contains("nvvm.reqntid = array"), + "{text}" + ); +} diff --git a/crates/cuda-oxide-mlir-export/tests/cute_gemv.rs b/crates/cuda-oxide-mlir-export/tests/cute_gemv.rs new file mode 100644 index 0000000000..9e42f6c4bb --- /dev/null +++ b/crates/cuda-oxide-mlir-export/tests/cute_gemv.rs @@ -0,0 +1,260 @@ +/* + * 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_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, +}; +use pliron_mlir_export::{MlirModule, render_module}; + +fn append(block: Ptr, operation: O, ctx: &Context) -> Ptr { + let operation = operation.get_operation(); + operation.insert_at_back(block, ctx); + operation +} + +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 +} + +fn translate_gemv_target(ctx: &Context, module: &ModuleOp) -> MlirModule { + CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(ctx, module) + .unwrap() +} + +fn translate_gemv(ctx: &Context, module: &ModuleOp) -> String { + render_module(&translate_gemv_target(ctx, module)) +} + +#[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/crates/mir-importer/src/lib.rs b/crates/mir-importer/src/lib.rs index 2909011287..c51af6c482 100644 --- a/crates/mir-importer/src/lib.rs +++ b/crates/mir-importer/src/lib.rs @@ -101,6 +101,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, diff --git a/crates/mir-importer/src/pipeline.rs b/crates/mir-importer/src/pipeline.rs index c1deef0a8f..e7d02fc248 100644 --- a/crates/mir-importer/src/pipeline.rs +++ b/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::{DeviceExternAttrs, DeviceExternDecl, PipelineError}; use llvm_export::export::DebugKind; @@ -189,6 +189,13 @@ pub struct PipelineConfig { pub output_dir: std::path::PathBuf, /// Base name for output files (e.g., `"kernel"` → `kernel.ll`, `kernel.ptx`). pub output_name: String, + /// Optional exact path for a textual CUTLASS MLIR observation. + /// + /// The export observes the prepared semantic module before the selected + /// backend continuation; failures are compilation errors. + pub mlir_output: Option, + /// Device artifact backend. The ordinary backend remains the default. + pub device_backend: DeviceBackend, /// Print progress messages to stdout. pub verbose: bool, /// Dump the `dialect-mir` module after translation (for debugging). @@ -253,6 +260,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, @@ -506,6 +515,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 { @@ -568,6 +578,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); @@ -577,6 +590,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); @@ -586,9 +602,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 { @@ -630,6 +653,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, + }), } } @@ -754,6 +786,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); @@ -1284,6 +1318,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, @@ -1337,6 +1373,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, @@ -1428,6 +1466,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/crates/pliron-mlir-export/Cargo.toml b/crates/pliron-mlir-export/Cargo.toml new file mode 100644 index 0000000000..c6f8fd25c2 --- /dev/null +++ b/crates/pliron-mlir-export/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "pliron-mlir-export" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +description = "Extensible Pliron to textual MLIR translation engine" +readme = "README.md" + +[dependencies] +pliron = { workspace = true } +thiserror = { workspace = true } + +[dev-dependencies] +pliron-derive = { workspace = true } diff --git a/crates/pliron-mlir-export/README.md b/crates/pliron-mlir-export/README.md new file mode 100644 index 0000000000..a1efb16bd4 --- /dev/null +++ b/crates/pliron-mlir-export/README.md @@ -0,0 +1,20 @@ +# Pliron to MLIR export + +This crate translates a Pliron module into textual MLIR. + +It deliberately knows nothing about CUDA, CuTe, or a particular MLIR build. +Those details are supplied as mapping packs: + +```text +Pliron module + │ + ▼ +mapping registry ──► small MLIR syntax tree ──► deterministic text + ▲ + │ +core / NVVM / CuTe mapping packs +``` + +Every source operation, type, and attribute needs an explicit mapping. A +matching name is never treated as permission to copy an unknown item through. +That makes dialect drift a clear translation error instead of malformed MLIR. diff --git a/crates/pliron-mlir-export/src/ast.rs b/crates/pliron-mlir-export/src/ast.rs new file mode 100644 index 0000000000..8b181f5c19 --- /dev/null +++ b/crates/pliron-mlir-export/src/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/crates/pliron-mlir-export/src/error.rs b/crates/pliron-mlir-export/src/error.rs new file mode 100644 index 0000000000..7ab8b384e5 --- /dev/null +++ b/crates/pliron-mlir-export/src/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/crates/pliron-mlir-export/src/lib.rs b/crates/pliron-mlir-export/src/lib.rs new file mode 100644 index 0000000000..fc2927633b --- /dev/null +++ b/crates/pliron-mlir-export/src/lib.rs @@ -0,0 +1,42 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Build textual MLIR from any Pliron dialect. +//! +//! The engine has three small pieces: +//! +//! ```text +//! Pliron values and blocks +//! │ explicit mappings +//! ▼ +//! typed MLIR syntax tree +//! │ one renderer +//! ▼ +//! deterministic generic MLIR +//! ``` +//! +//! A mapping pack registers the operations, types, and attributes it owns. +//! Unknown items are collected into one error before translation starts. +//! Equal-looking source and target names are never copied implicitly. + +mod ast; +mod error; +mod registry; +mod render; +mod translate; + +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 render::render_module; +pub use translate::{ + LocationMode, OperationInput, TranslationConfig, TranslationSession, translate_module, +}; diff --git a/crates/pliron-mlir-export/src/registry.rs b/crates/pliron-mlir-export/src/registry.rs new file mode 100644 index 0000000000..9fe3a0f3be --- /dev/null +++ b/crates/pliron-mlir-export/src/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::translate::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/crates/pliron-mlir-export/src/render.rs b/crates/pliron-mlir-export/src/render.rs new file mode 100644 index 0000000000..230d06554c --- /dev/null +++ b/crates/pliron-mlir-export/src/render.rs @@ -0,0 +1,472 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use std::fmt::Write; + +use crate::{ + MlirAttribute, MlirBlock, MlirFloatType, MlirLocation, MlirModule, MlirOperation, MlirRegion, + MlirType, +}; + +/// Render one module using MLIR's generic operation syntax. +pub fn render_module(module: &MlirModule) -> String { + let mut output = String::new(); + render_operation(&mut output, &module.root, 0); + output +} + +fn render_operation(output: &mut String, operation: &MlirOperation, indent: usize) { + write_indent(output, indent); + if !operation.results.is_empty() { + for (index, result) in operation.results.iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + let _ = write!(output, "%v{}", result.id.0); + } + output.push_str(" = "); + } + let _ = write!(output, "\"{}\"(", operation.name); + for (index, operand) in operation.operands.iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + let _ = write!(output, "%v{}", operand.id.0); + } + output.push(')'); + + if !operation.successors.is_empty() { + output.push_str(" ["); + for (index, successor) in operation.successors.iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + let _ = write!(output, "^bb{}", successor.0); + } + output.push(']'); + } + + if !operation.properties.is_empty() { + output.push_str(" <{"); + render_named_attributes(output, &operation.properties); + output.push_str("}>"); + } + + if !operation.regions.is_empty() { + output.push_str(" ("); + for (index, region) in operation.regions.iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + render_region(output, region, indent); + } + output.push(')'); + } + + if !operation.attributes.is_empty() { + output.push_str(" {"); + render_named_attributes(output, &operation.attributes); + output.push('}'); + } + + output.push_str(" : ("); + for (index, operand) in operation.operands.iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + render_type(output, &operand.ty); + } + output.push_str(") -> "); + render_result_types(output, operation); + if operation.location != MlirLocation::Unknown { + output.push(' '); + render_location(output, &operation.location); + } + output.push('\n'); +} + +fn render_named_attributes( + output: &mut String, + attributes: &std::collections::BTreeMap, +) { + for (index, (name, attribute)) in attributes.iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + let _ = write!(output, "{name} = "); + render_attribute(output, attribute); + } +} + +fn render_result_types(output: &mut String, operation: &MlirOperation) { + match operation.results.as_slice() { + [] => output.push_str("()"), + [result] => render_type(output, &result.ty), + results => { + output.push('('); + for (index, result) in results.iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + render_type(output, &result.ty); + } + output.push(')'); + } + } +} + +fn render_region(output: &mut String, region: &MlirRegion, indent: usize) { + output.push_str("{\n"); + for block in ®ion.blocks { + render_block(output, block, indent + 1); + } + write_indent(output, indent); + output.push('}'); +} + +fn render_block(output: &mut String, block: &MlirBlock, indent: usize) { + write_indent(output, indent); + let _ = write!(output, "^bb{}", block.id.0); + if !block.arguments.is_empty() { + output.push('('); + for (index, argument) in block.arguments.iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + let _ = write!(output, "%v{}: ", argument.id.0); + render_type(output, &argument.ty); + if argument.location != MlirLocation::Unknown { + output.push(' '); + render_location(output, &argument.location); + } + } + output.push(')'); + } + output.push_str(":\n"); + for operation in &block.operations { + render_operation(output, operation, indent + 1); + } +} + +fn render_type(output: &mut String, ty: &MlirType) { + match ty { + MlirType::Index => output.push_str("index"), + MlirType::Integer(width) => { + let _ = write!(output, "i{width}"); + } + MlirType::Float(kind) => output.push_str(match kind { + MlirFloatType::F16 => "f16", + MlirFloatType::BF16 => "bf16", + MlirFloatType::F32 => "f32", + MlirFloatType::F64 => "f64", + MlirFloatType::Other(spelling) => spelling, + }), + MlirType::Vector { shape, element } => { + output.push_str("vector<"); + for dimension in shape { + let _ = write!(output, "{dimension}x"); + } + render_type(output, element); + output.push('>'); + } + MlirType::Function { inputs, results } => { + render_type_list(output, inputs); + output.push_str(" -> "); + if results.len() == 1 { + render_type(output, &results[0]); + } else { + render_type_list(output, results); + } + } + MlirType::Tuple(elements) => { + output.push_str("tuple<"); + for (index, element) in elements.iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + render_type(output, element); + } + output.push('>'); + } + MlirType::Dialect(spelling) => output.push_str(spelling), + } +} + +fn render_type_list(output: &mut String, types: &[MlirType]) { + output.push('('); + for (index, ty) in types.iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + render_type(output, ty); + } + output.push(')'); +} + +fn render_attribute(output: &mut String, attribute: &MlirAttribute) { + match attribute { + MlirAttribute::Unit => output.push_str("unit"), + MlirAttribute::Bool(value) => output.push_str(if *value { "true" } else { "false" }), + MlirAttribute::Integer { value, ty } => { + let _ = write!(output, "{value} : "); + render_type(output, ty); + } + MlirAttribute::Float { value, ty } => { + output.push_str(value); + output.push_str(" : "); + render_type(output, ty); + } + MlirAttribute::String(value) => render_string(output, value), + MlirAttribute::Type(ty) => render_type(output, ty), + MlirAttribute::Array(elements) => { + output.push('['); + for (index, element) in elements.iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + render_attribute(output, element); + } + output.push(']'); + } + MlirAttribute::DenseI32Array(elements) => { + output.push_str("array'); + } + MlirAttribute::DenseI64Array(elements) => { + output.push_str("array'); + } + MlirAttribute::Dictionary(entries) => { + output.push('{'); + for (index, (name, value)) in entries.iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + let _ = write!(output, "{name} = "); + render_attribute(output, value); + } + output.push('}'); + } + MlirAttribute::SymbolRef(symbol) => { + output.push('@'); + output.push_str(symbol); + } + MlirAttribute::Dialect(spelling) => output.push_str(spelling), + } +} + +fn render_location(output: &mut String, location: &MlirLocation) { + output.push_str("loc("); + render_location_instance(output, location); + output.push(')'); +} + +/// A bare location instance, as MLIR's grammar nests them: only the +/// outermost location carries the `loc(...)` wrapper, so a named or fused +/// child is spelled `"name"(unknown)`, never `"name"(loc(unknown))`. +fn render_location_instance(output: &mut String, location: &MlirLocation) { + match location { + MlirLocation::Unknown => output.push_str("unknown"), + MlirLocation::FileLineCol { file, line, column } => { + render_string(output, file); + let _ = write!(output, ":{line}:{column}"); + } + MlirLocation::Fused { + metadata, + locations, + } => { + output.push_str("fused"); + if let Some(metadata) = metadata { + output.push('<'); + render_attribute(output, metadata); + output.push('>'); + } + output.push('['); + for (index, child) in locations.iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + render_location_instance(output, child); + } + output.push(']'); + } + MlirLocation::Named { name, child } => { + render_string(output, name); + // MLIR elides an unknown child: `loc("name")`. + if !matches!(**child, MlirLocation::Unknown) { + output.push('('); + render_location_instance(output, child); + output.push(')'); + } + } + MlirLocation::CallSite { callee, caller } => { + output.push_str("callsite("); + render_location_instance(output, callee); + output.push_str(" at "); + render_location_instance(output, caller); + output.push(')'); + } + } +} + +fn render_string(output: &mut String, value: &str) { + output.push('"'); + for character in value.chars() { + match character { + '"' => output.push_str("\\\""), + '\\' => output.push_str("\\\\"), + '\n' => output.push_str("\\n"), + '\r' => output.push_str("\\r"), + '\t' => output.push_str("\\t"), + character => output.push(character), + } + } + output.push('"'); +} + +fn write_indent(output: &mut String, indent: usize) { + for _ in 0..indent { + output.push_str(" "); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{MlirBlockArgument, MlirBlockId, MlirResult, MlirValueId, MlirValueUse}; + + #[test] + fn generic_syntax_is_stable() { + let i32_ty = MlirType::Integer(32); + let add = MlirOperation { + name: "arith.addi".into(), + results: vec![MlirResult { + id: MlirValueId(2), + ty: i32_ty.clone(), + }], + operands: vec![ + MlirValueUse { + id: MlirValueId(0), + ty: i32_ty.clone(), + }, + MlirValueUse { + id: MlirValueId(1), + ty: i32_ty.clone(), + }, + ], + successors: vec![], + properties: Default::default(), + regions: vec![], + attributes: Default::default(), + location: MlirLocation::Unknown, + }; + let mut root = MlirOperation::new("builtin.module").unwrap(); + root.regions.push(MlirRegion { + blocks: vec![MlirBlock { + id: MlirBlockId(0), + arguments: vec![ + MlirBlockArgument { + id: MlirValueId(0), + ty: i32_ty.clone(), + location: MlirLocation::Unknown, + }, + MlirBlockArgument { + id: MlirValueId(1), + ty: i32_ty, + location: MlirLocation::Unknown, + }, + ], + operations: vec![add], + }], + }); + let text = render_module(&MlirModule { + root, + profile: "test".into(), + }); + assert_eq!( + text, + "\"builtin.module\"() ({\n ^bb0(%v0: i32, %v1: i32):\n %v2 = \"arith.addi\"(%v0, %v1) : (i32, i32) -> i32\n}) : () -> ()\n" + ); + } + + #[test] + fn operation_properties_use_mlir_generic_syntax() { + let i1_ty = MlirType::Integer(1); + let i32_ty = MlirType::Integer(32); + let mut branch = MlirOperation::new("cf.cond_br").unwrap(); + branch.operands = vec![ + MlirValueUse { + id: MlirValueId(0), + ty: i1_ty.clone(), + }, + MlirValueUse { + id: MlirValueId(1), + ty: i32_ty.clone(), + }, + MlirValueUse { + id: MlirValueId(1), + ty: i32_ty, + }, + ]; + branch.successors = vec![MlirBlockId(1), MlirBlockId(2)]; + branch.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 1, 1]), + ); + + let mut root = MlirOperation::new("builtin.module").unwrap(); + root.regions.push(MlirRegion { + blocks: vec![MlirBlock { + id: MlirBlockId(0), + arguments: vec![ + MlirBlockArgument { + id: MlirValueId(0), + ty: i1_ty, + location: MlirLocation::Unknown, + }, + MlirBlockArgument { + id: MlirValueId(1), + ty: MlirType::Integer(32), + location: MlirLocation::Unknown, + }, + ], + operations: vec![branch], + }], + }); + let text = render_module(&MlirModule { + root, + profile: "test".into(), + }); + + assert!(text.contains( + "\"cf.cond_br\"(%v0, %v1, %v1) [^bb1, ^bb2] <{operandSegmentSizes = array}> : (i1, i32, i32) -> ()" + )); + } + + #[test] + fn dense_i64_properties_render_without_stringly_typed_escapes() { + let mut text = String::new(); + render_attribute(&mut text, &MlirAttribute::DenseI64Array(vec![0, 2])); + assert_eq!(text, "array"); + } +} diff --git a/crates/pliron-mlir-export/src/translate.rs b/crates/pliron-mlir-export/src/translate.rs new file mode 100644 index 0000000000..2ca0a0c044 --- /dev/null +++ b/crates/pliron-mlir-export/src/translate.rs @@ -0,0 +1,839 @@ +/* + * 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::{Located, 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(()) +} + +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, + 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); + assert!(text.starts_with("\"builtin.module\"() ({")); + assert!(text.ends_with(": () -> ()\n")); + } + + #[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), crate::render_module(&second)); + } + + #[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/crates/rustc-codegen-cuda/Cargo.lock b/crates/rustc-codegen-cuda/Cargo.lock index 260533a397..1f593cc64c 100644 --- a/crates/rustc-codegen-cuda/Cargo.lock +++ b/crates/rustc-codegen-cuda/Cargo.lock @@ -205,24 +205,45 @@ 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", + "pliron-mlir-export", "ptx-parse", "rustc-demangle", + "serde_json", + "sha2", "tempfile", "thiserror", ] +[[package]] +name = "cuda-oxide-mlir-export" +version = "0.1.0" +dependencies = [ + "dialect-cute", + "dialect-mir", + "dialect-nvvm", + "llvm-export", + "pliron", + "pliron-mlir-export", + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -662,6 +683,14 @@ dependencies = [ "thiserror", ] +[[package]] +name = "pliron-mlir-export" +version = "0.1.0" +dependencies = [ + "pliron", + "thiserror", +] + [[package]] name = "proc-macro2" version = "1.0.106" diff --git a/crates/rustc-codegen-cuda/README.md b/crates/rustc-codegen-cuda/README.md index ed36502cee..cb31611ffc 100644 --- a/crates/rustc-codegen-cuda/README.md +++ b/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) @@ -87,6 +89,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 +126,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 diff --git a/crates/rustc-codegen-cuda/src/device_codegen.rs b/crates/rustc-codegen-cuda/src/device_codegen.rs index d9868c1571..30369965d8 100644 --- a/crates/rustc-codegen-cuda/src/device_codegen.rs +++ b/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, @@ -997,6 +1003,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; @@ -1203,6 +1211,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, @@ -1364,6 +1374,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/crates/rustc-codegen-cuda/src/lib.rs b/crates/rustc-codegen-cuda/src/lib.rs index 8deaf0cf0f..aee9c9be9f 100644 --- a/crates/rustc-codegen-cuda/src/lib.rs +++ b/crates/rustc-codegen-cuda/src/lib.rs @@ -285,6 +285,8 @@ //! | `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_TARGET` | Override GPU target (e.g., `sm_90a`) | //! | `CUDA_OXIDE_DEVICE_CODEGEN_CRATE` | Filter device owner crate names | //! @@ -374,7 +376,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 +388,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 +407,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 +421,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 +438,76 @@ 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()); + } + }; + parse_device_backend( + selector.as_deref(), + std::env::var_os("CUDA_OXIDE_CUTLASS_COMPILER").map(std::path::PathBuf::from), + profile.as_deref(), + ) +} + +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 +604,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 +702,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 +753,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, @@ -1236,6 +1333,62 @@ 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 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-book/getting-started/installation.md b/cuda-oxide-book/getting-started/installation.md index 907010f59e..196f6ebc6f 100644 --- a/cuda-oxide-book/getting-started/installation.md +++ b/cuda-oxide-book/getting-started/installation.md @@ -165,6 +165,29 @@ 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. + +--- + ## LLVM 21+ (optional) cuda-oxide uses LLVM's NVPTX backend to lower LLVM IR to PTX. @@ -321,6 +344,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: diff --git a/cuteir/README.md b/cuteir/README.md index fa17f4d7e4..038f41a649 100644 --- a/cuteir/README.md +++ b/cuteir/README.md @@ -87,3 +87,8 @@ own READMEs give the exact device and shape requirements. Each program checks every GPU output bit pattern against a deterministic host result before reporting success. Backend-specific artifact and code-shape checks live with the backend implementation rather than in this shared layer. + +On the translation branch, the [official CUTLASS backend](docs/translation.md) +maps these semantic operations directly to the CUTLASS 4.7 MLIR profile and +embeds the compiler-produced cubin. It does not pass through native CuTe +expansion. diff --git a/cuteir/docs/translation.md b/cuteir/docs/translation.md new file mode 100644 index 0000000000..5c37d830b4 --- /dev/null +++ b/cuteir/docs/translation.md @@ -0,0 +1,77 @@ +# Official CUTLASS translation backend + +The translation backend preserves the high-level `dialect-cute` program after +common preparation, verifies the complete semantic story, and maps it directly +to the MLIR dialects accepted by NVIDIA's official CUTLASS 4.7 compiler. + +```text +cute-rs kernel + │ + ▼ +high-level dialect-cute + │ + ├── common preparation + └── backend-neutral whole-module verification + │ + ▼ +CUTLASS full-CuTe MLIR profile + │ + ▼ +official libCutlassCompiler 4.7 + │ + ▼ +validated cubin → embedded host artifact → CUDA launch +``` + +The backend contains direct mapping packs for all three shared examples: + +- tensor tiling and copying for `elementwise_cute`; +- scaled tensor views and NVFP4 GEMV for `nvfp4_gemv_cute`; +- scheduler, work tiles, TMA load pipelines, shared-memory MMA, and epilogue + stores for `blockscale_gemm_cute`. + +It does not run a native CuTe expansion pass or lower through backend-specific +MIR/NVVM leaf intrinsics. In the GEMM epilogue, +`ReadyForTma` directly emits the async-shared proxy publication fence followed +by the counted CTA hand-off barrier. `Reusable` emits only the counted barrier. + +## Install and select the compiler + +Install the pinned compiler archive once: + +```bash +cargo oxide toolchain install cutlass +``` + +The installer verifies the official archive and `libCutlassCompiler.so` +digests. With no explicit compiler path, `cargo oxide` resolves this managed +installation and includes its library digest in the build fingerprint. + +Select the backend for a build: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide build elementwise_cute --arch sm_120a + +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide build nvfp4_gemv_cute --arch sm_120a + +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide build blockscale_gemm_cute --arch sm_120a +``` + +An explicit absolute library path overrides the managed installation: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ +CUDA_OXIDE_CUTLASS_COMPILER=/absolute/path/to/libCutlassCompiler.so \ + cargo oxide build elementwise_cute --arch sm_120a +``` + +Use `CUDA_OXIDE_MLIR_OUTPUT=/absolute/path/module.mlir` for an optional textual +observation. MLIR dumps and generated cubins are build artifacts and must not be +committed. + +The ordinary CUDA Oxide MIR/NVVM/LLVM-to-PTX backend remains the default for +non-CuTe kernels. Selecting the CUTLASS backend changes only the continuation +after shared preparation and semantic verification. diff --git a/cuteir/examples/blockscale_gemm_cute/Cargo.lock b/cuteir/examples/blockscale_gemm_cute/Cargo.lock index 6a848bbf7f..cd0219bd05 100644 --- a/cuteir/examples/blockscale_gemm_cute/Cargo.lock +++ b/cuteir/examples/blockscale_gemm_cute/Cargo.lock @@ -181,6 +181,10 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + [[package]] name = "cute-layout" version = "0.1.0" @@ -303,6 +307,7 @@ dependencies = [ name = "libnvvm-sys" version = "0.2.1" dependencies = [ + "cuda-target-spec", "libloading", "thiserror", ] diff --git a/cuteir/examples/elementwise_cute/Cargo.lock b/cuteir/examples/elementwise_cute/Cargo.lock index f8ab63ab9f..3d7c4c6c21 100644 --- a/cuteir/examples/elementwise_cute/Cargo.lock +++ b/cuteir/examples/elementwise_cute/Cargo.lock @@ -171,6 +171,10 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + [[package]] name = "cute-layout" version = "0.1.0" @@ -302,6 +306,7 @@ dependencies = [ name = "libnvvm-sys" version = "0.2.1" dependencies = [ + "cuda-target-spec", "libloading", "thiserror", ] diff --git a/cuteir/examples/nvfp4_gemv_cute/Cargo.lock b/cuteir/examples/nvfp4_gemv_cute/Cargo.lock index 9394caaae8..567f3f0f2b 100644 --- a/cuteir/examples/nvfp4_gemv_cute/Cargo.lock +++ b/cuteir/examples/nvfp4_gemv_cute/Cargo.lock @@ -171,6 +171,10 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + [[package]] name = "cute-layout" version = "0.1.0" @@ -292,6 +296,7 @@ dependencies = [ name = "libnvvm-sys" version = "0.2.1" dependencies = [ + "cuda-target-spec", "libloading", "thiserror", ] diff --git a/dependency-licenses.csv b/dependency-licenses.csv index afd2e56ef4..68ef54d01d 100644 --- a/dependency-licenses.csv +++ b/dependency-licenses.csv @@ -35,6 +35,7 @@ cuda-intrinsics,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),htt cuda-intrinsics-gen,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,Extractor and deterministic source generator for cuda-oxide intrinsics, cuda-macros,0.2.1,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,"Procedural macros for CUDA kernel development (#[kernel], #[device], etc.)", cuda-oxide-codegen,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,Rustc-independent cuda-oxide PTX backend, +cuda-oxide-mlir-export,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,CUDA Oxide mapping packs for the official CUTLASS MLIR consumer profile, cuda-target-spec,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,Shared CUDA target parsing and recorded LLVM PTX-floor policy, cute-layout,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,Coordinate-to-memory mapping and validation for CuTe-style GPU tiles, cute-rs,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,"CuTe-style device library: Tensor, Layout types, partitioning, copy/mma atoms", @@ -83,6 +84,7 @@ oxide-artifacts,0.2.1,Apache-2.0,Yes,registry+https://github.com/rust-lang/crate pliron,0.17.0,Apache-2.0,Yes,git+https://github.com/pliron-org/pliron?rev=8447aa426e2f090dc8b7b9383c5035db7bc70b62,3rd Party (git),https://github.com/pliron-org/pliron,Programming Languages Intermediate RepresentatiON,"Primary dependency: MLIR-inspired IR framework. Custom MLIR dialects (MIR, LLVM, NVVM) built on this. Pinned to a git revision; this version is not published to crates.io." pliron-derive,0.17.0,Apache-2.0,Yes,git+https://github.com/pliron-org/pliron?rev=8447aa426e2f090dc8b7b9383c5035db7bc70b62,3rd Party (git),https://github.com/pliron-org/pliron,LLVM Dialect for pliron,"Proc-macro companion to pliron: the Description column copies upstream's own text, which is inaccurate for this crate. Same git pin as pliron." pliron-llvm,0.17.0,Apache-2.0,Yes,git+https://github.com/pliron-org/pliron?rev=8447aa426e2f090dc8b7b9383c5035db7bc70b62,3rd Party (git),https://github.com/pliron-org/pliron,LLVM dialect for pliron,"Upstream LLVM dialect modeling; re-exported by the llvm-export crate. Same git pin as pliron." +pliron-mlir-export,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,Typed deterministic translation from Pliron modules to textual MLIR, prettyplease,0.2.37,MIT OR Apache-2.0,Yes,registry+https://github.com/rust-lang/crates.io-index,3rd Party (crates.io),https://github.com/dtolnay/prettyplease,A minimal `syn` syntax tree pretty-printer, proc-macro2,1.0.106,MIT OR Apache-2.0,Yes,registry+https://github.com/rust-lang/crates.io-index,3rd Party (crates.io),https://github.com/dtolnay/proc-macro2,A substitute implementation of the compiler's `proc_macro` API to decouple token-based libraries from the procedural macro use case., ptx-parse,0.1.0,Apache-2.0,Yes,local (workspace),1st Party (workspace),https://github.com/NVlabs/cuda-oxide.git,Lossless structural views over PTX source text, From f33f2413209f94b7ec1ef0cac4b51e9ad019d041 Mon Sep 17 00:00:00 2001 From: nihalp Date: Tue, 18 Aug 2026 09:26:38 +0530 Subject: [PATCH 03/99] docs(cuteir): record CUTLASS performance comparison --- .../examples/blockscale_gemm_cute/README.md | 39 +++++++++++++++-- cuteir/examples/nvfp4_gemv_cute/README.md | 42 +++++++++++++++++-- 2 files changed, 74 insertions(+), 7 deletions(-) diff --git a/cuteir/examples/blockscale_gemm_cute/README.md b/cuteir/examples/blockscale_gemm_cute/README.md index 1705edb4eb..feee4a516e 100644 --- a/cuteir/examples/blockscale_gemm_cute/README.md +++ b/cuteir/examples/blockscale_gemm_cute/README.md @@ -94,12 +94,23 @@ target `sm_120a`. On another GPU, the program skips device execution. Numerical GEMM validation requires a compatible SM 12.x GPU. -## Run it +## Build and run it -From the repository root: +From the repository root, install the pinned official CUTLASS 4.7 compiler +once: ```bash -cargo oxide run blockscale_gemm_cute --arch sm_120a +cargo oxide toolchain install cutlass +``` + +Build or run this example through the translation backend: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide build blockscale_gemm_cute --arch sm_120a + +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide run blockscale_gemm_cute --arch sm_120a ``` The program builds deterministic packed inputs, computes a host reference, @@ -112,6 +123,28 @@ then performs 100 warmup launches and measures 31 individual launches with CUDA events. The report includes the median, p10, and p90 device time and the median-derived TFLOP/s for `2 × M × N × K` floating-point operations. +## 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=N=2048`, +`K=1024`, `L=1`: + +| Implementation | Median | Median-derived rate | +| :--- | ---: | ---: | +| CUTLASS 4.7 translation | **11.264 µs** | 762.60 TFLOP/s | +| CuTeDSL 4.6.2 | **10.272 µs** | 836.25 TFLOP/s | + +Translation is 0.992 µs, or 9.66%, longer. The measured translation cubin is +29,608 bytes with SHA-256 +`69e1f1aad53a5889a40f6d7bfce0c6f66607eaeb1119640aa94f8ed3467dfbee`. +Translation matched all 4,194,304 expected `f16` bit patterns; the separate +CuTeDSL reference-checking run also passed. + +This is a direct hot-reuse comparison with serial submission and no CUDA +graphs. Both implementations ran in one process with 100 warmups and 31 +one-kernel samples; translation also performed one correctness launch. Nsight +measured only main-kernel active time, excluding setup and conversion kernels. + This is the full semantic acceptance example: ```text diff --git a/cuteir/examples/nvfp4_gemv_cute/README.md b/cuteir/examples/nvfp4_gemv_cute/README.md index 3f333336b9..a11786f487 100644 --- a/cuteir/examples/nvfp4_gemv_cute/README.md +++ b/cuteir/examples/nvfp4_gemv_cute/README.md @@ -86,13 +86,24 @@ L > 0 `L` is the batch count. The default problem is `M=512`, `K=256`, `L=1`. -## Run it +## Build and run it -From the repository root: +From the repository root, install the pinned official CUTLASS 4.7 compiler +once: ```bash -cargo oxide run nvfp4_gemv_cute --arch sm_120a -- \ - --m 512 --k 256 --l 1 +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 @@ -101,3 +112,26 @@ 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. From 68f2d1b6425660e96b2c70fee912ca57dffb6a26 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Wed, 2 Sep 2026 16:42:49 +0530 Subject: [PATCH 04/99] test(cute-rs): refresh trybuild snapshot for the pinned nightly's diagnostic rendering nightly-2026-08-28 prints the mismatched 'found' type in tiled_copy_role_mismatch as `SharedTensor<_, f16, (), Nkl>` where the snapshot (recorded under the previous nightly) had `SharedTensor<..., f16, (), ...>`. Same error, same span; re-blessed with TRYBUILD=overwrite. --- .../cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr b/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr index eadfa5e577..98337b9f4b 100644 --- a/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr +++ b/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr @@ -2,7 +2,7 @@ error[E0308]: mismatched types --> tests/compile_fail/tiled_copy_role_mismatch.rs:22:41 | 22 | unsafe { copy.copy(&source, (0, 0), &mut shared_b, 0) }; - | ---- ^^^^^^^^^^^^^ expected `&mut SharedTensor`, found `&mut SharedTensor<..., f16, (), ...>` + | ---- ^^^^^^^^^^^^^ expected `&mut SharedTensor`, found `&mut SharedTensor<_, f16, (), Nkl>` | | | arguments to this method are incorrect | From 60ff1edcb57b326c70c3f6f3e458eecbd613f28d Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Wed, 2 Sep 2026 20:59:47 +0530 Subject: [PATCH 05/99] chore(vscode): track the project rust-analyzer settings Same change as on main, with the three standalone cuteir example workspaces listed so rust-analyzer loads them on this branch. --- .vscode/settings.json | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index 65be775e00..12d1042dc6 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -11,13 +11,19 @@ "files.exclude": { "**/target/flycheck*": true, }, - // Every example under crates/rustc-codegen-cuda/examples/ is its own cargo - // workspace, so rust-analyzer only sees the ones listed here. Add the ones - // you are working on. Every path must exist on the checked-out branch, or - // rust-analyzer reports "Failed to load workspaces". + // Every example under crates/rustc-codegen-cuda/examples/ and cuteir/examples/ + // is its own cargo workspace, so rust-analyzer only sees the ones listed + // here. Add the ones you are working on. Every path must exist on the + // checked-out branch, or rust-analyzer reports "Failed to load workspaces". "rust-analyzer.linkedProjects": [ "./Cargo.toml", "./crates/rustc-codegen-cuda/Cargo.toml", + // cute-layout, cute-rs and dialect-cute are root-workspace members + // (covered by ./Cargo.toml above); there is no cuteir/Cargo.toml. + // The standalone cuteir examples need their own entries: + "./cuteir/examples/elementwise_cute/Cargo.toml", + "./cuteir/examples/blockscale_gemm_cute/Cargo.toml", + "./cuteir/examples/nvfp4_gemv_cute/Cargo.toml", "./crates/rustc-codegen-cuda/examples/debug_pointer_locals/Cargo.toml", "./crates/rustc-codegen-cuda/examples/gemm_sol_final/Cargo.toml", ], From bb1825599eee0b33252b86f906cc140e1fa670bb Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Wed, 16 Sep 2026 13:52:37 +0530 Subject: [PATCH 06/99] feat(cuteir): add SM100 CuTe GEMM and compiler support Express the persistent 256x352 FP16 GEMM with typed CuTe layouts, two-CTA paired MMA and collector reuse, cluster TMA, tensor memory, and pipeline operations through the CUTLASS translation backend. Preserve vectors through epilogue conversion and shared-memory stores. Group CUTLASS compiler regressions under src/tests/cutlass, add GPU validation cases, and document the example with concise diagrams and single-launch commands. Keep the local Python reference, comparison harness, benchmark results, and performance notes out of version control. Validation: pinned CUTLASS compiler regressions, exporter/frontend checks, six GPU correctness cases, memcheck, synccheck, and single-launch Python oracle comparisons for both performance shapes passed. --- .../cuda-oxide-codegen/src/cutlass_backend.rs | 8 + crates/cuda-oxide-codegen/src/prep.rs | 44 + .../src/tests/cutlass/sm100.rs | 197 ++++ .../src/tests/cutlass/tcgen05.rs | 348 ++++++ crates/cuda-oxide-mlir-export/README.md | 8 + .../cuda-oxide-mlir-export/src/cute_sm100.rs | 894 ++++++++++++++ crates/cuda-oxide-mlir-export/src/lib.rs | 6 + crates/cuda-oxide-mlir-export/src/mir_core.rs | 186 ++- .../cuda-oxide-mlir-export/src/mir_memory.rs | 129 ++- .../src/nvvm_cluster.rs | 1029 +++++++++++++++++ .../src/nvvm_tcgen05.rs | 598 ++++++++++ crates/cuda-oxide-mlir-export/src/profile.rs | 6 +- .../tests/cute_elementwise.rs | 14 +- .../src/translator/cute/layout.rs | 6 +- .../mir-importer/src/translator/cute/mod.rs | 109 ++ .../src/translator/cute/sm100_emit.rs | 358 ++++++ cuteir/README.md | 7 +- cuteir/cute-rs/src/lib.rs | 5 + cuteir/cute-rs/src/sm100.rs | 491 ++++++++ .../compile_fail/sm100_mma_shape_mismatch.rs | 9 + .../sm100_mma_shape_mismatch.stderr | 11 + .../sm100_operand_role_mismatch.rs | 16 + .../sm100_operand_role_mismatch.stderr | 15 + .../compile_fail/sm100_tma_layout_mismatch.rs | 15 + .../sm100_tma_layout_mismatch.stderr | 15 + cuteir/cute-rs/tests/type_safety.rs | 3 + cuteir/dialect-cute/src/lib.rs | 8 +- cuteir/dialect-cute/src/sm100_ops.rs | 913 +++++++++++++++ cuteir/dialect-cute/src/verify.rs | 6 +- .../docs/assets/cutlass-translation-flow.svg | 47 + cuteir/docs/translation.md | 130 ++- .../fp16_gemm_256x352_cute/.gitignore | 13 + .../fp16_gemm_256x352_cute/Cargo.lock | 598 ++++++++++ .../fp16_gemm_256x352_cute/Cargo.toml | 15 + .../examples/fp16_gemm_256x352_cute/README.md | 116 ++ .../fp16_gemm_256x352_cute/src/kernel.rs | 297 +++++ .../fp16_gemm_256x352_cute/src/main.rs | 559 +++++++++ .../fp16_gemm_256x352_cute/validate.py | 161 +++ 38 files changed, 7312 insertions(+), 78 deletions(-) create mode 100644 crates/cuda-oxide-codegen/src/tests/cutlass/sm100.rs create mode 100644 crates/cuda-oxide-codegen/src/tests/cutlass/tcgen05.rs create mode 100644 crates/cuda-oxide-mlir-export/src/cute_sm100.rs create mode 100644 crates/cuda-oxide-mlir-export/src/nvvm_cluster.rs create mode 100644 crates/cuda-oxide-mlir-export/src/nvvm_tcgen05.rs create mode 100644 crates/mir-importer/src/translator/cute/sm100_emit.rs create mode 100644 cuteir/cute-rs/src/sm100.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.stderr create mode 100644 cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.stderr create mode 100644 cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.rs create mode 100644 cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.stderr create mode 100644 cuteir/dialect-cute/src/sm100_ops.rs create mode 100644 cuteir/docs/assets/cutlass-translation-flow.svg create mode 100644 cuteir/examples/fp16_gemm_256x352_cute/.gitignore create mode 100644 cuteir/examples/fp16_gemm_256x352_cute/Cargo.lock create mode 100644 cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml create mode 100644 cuteir/examples/fp16_gemm_256x352_cute/README.md create mode 100644 cuteir/examples/fp16_gemm_256x352_cute/src/kernel.rs create mode 100644 cuteir/examples/fp16_gemm_256x352_cute/src/main.rs create mode 100644 cuteir/examples/fp16_gemm_256x352_cute/validate.py diff --git a/crates/cuda-oxide-codegen/src/cutlass_backend.rs b/crates/cuda-oxide-codegen/src/cutlass_backend.rs index fd5c2900ed..57b399abd6 100644 --- a/crates/cuda-oxide-codegen/src/cutlass_backend.rs +++ b/crates/cuda-oxide-codegen/src/cutlass_backend.rs @@ -39,6 +39,14 @@ const CUTLASS_EXTERNAL_MODULE_ABI: &str = /// 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; + /// External input needed to execute the official CUTLASS Backend B. #[derive(Clone, Debug, PartialEq, Eq)] pub struct CutlassBackendConfig { diff --git a/crates/cuda-oxide-codegen/src/prep.rs b/crates/cuda-oxide-codegen/src/prep.rs index 998559ab87..81f642f2d5 100644 --- a/crates/cuda-oxide-codegen/src/prep.rs +++ b/crates/cuda-oxide-codegen/src/prep.rs @@ -26,6 +26,7 @@ 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, @@ -290,6 +291,25 @@ fn has_high_level_cute_operations(ctx: &Context, root: Ptr) -> bool { || 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; } @@ -762,6 +782,30 @@ mod tests { assert!(has_high_level_cute_operations(&ctx, module_operation)); } + #[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::(); diff --git a/crates/cuda-oxide-codegen/src/tests/cutlass/sm100.rs b/crates/cuda-oxide-codegen/src/tests/cutlass/sm100.rs new file mode 100644 index 0000000000..41c35c3114 --- /dev/null +++ b/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 = pliron_mlir_export::render_module(&target); + 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/crates/cuda-oxide-codegen/src/tests/cutlass/tcgen05.rs b/crates/cuda-oxide-codegen/src/tests/cutlass/tcgen05.rs new file mode 100644 index 0000000000..5936d30e9a --- /dev/null +++ b/crates/cuda-oxide-codegen/src/tests/cutlass/tcgen05.rs @@ -0,0 +1,348 @@ +/* + * 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, + FenceMbarrierInitReleaseClusterOp, FenceProxyAsyncSharedCtaOp, MapaSharedClusterOp, + MbarrierArriveClusterOp, MbarrierArriveExpectTxClusterOp, MbarrierInitSharedOp, + MbarrierTryWaitParityClusterOp, 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(&mut ctx, Tcgen05MmaFormAttr::Shared); + mma.set_attr_nvvm_tcgen05_mma_kind(&mut ctx, Tcgen05MmaKindAttr::F16); + mma.set_attr_nvvm_tcgen05_mma_cta_group(&mut ctx, Tcgen05MmaCtaGroupAttr::Cg2); + mma.set_attr_nvvm_tcgen05_mma_collector_a(&mut 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 = pliron_mlir_export::render_module(&target); + 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(&mut 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(&mut 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![], vec![], None); + + let profile = CutlassFullCuteMlir22::new("sm_100a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + let mlir = pliron_mlir_export::render_module(&target); + 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!("Cluster/TMA translated MLIR failed compilation: {error}\n{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) +} diff --git a/crates/cuda-oxide-mlir-export/README.md b/crates/cuda-oxide-mlir-export/README.md index b015cbddcb..5299d09e55 100644 --- a/crates/cuda-oxide-mlir-export/README.md +++ b/crates/cuda-oxide-mlir-export/README.md @@ -42,6 +42,14 @@ artifact loader. Set `CUDA_OXIDE_MLIR_OUTPUT=` to retain the exact textual module for inspection; 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`](../../cuteir/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. + ## The first mapping pack The first pack covers the ordinary scalar code around a CuTe kernel: diff --git a/crates/cuda-oxide-mlir-export/src/cute_sm100.rs b/crates/cuda-oxide-mlir-export/src/cute_sm100.rs new file mode 100644 index 0000000000..a4fe9c30bd --- /dev/null +++ b/crates/cuda-oxide-mlir-export/src/cute_sm100.rs @@ -0,0 +1,894 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SM100 CuTe semantic plans lowered by the pinned CUTLASS 4.7 compiler. +//! +//! Allocation, tensor-memory copies, cluster TMA, and shared-memory descriptors +//! use native CuTe operations. The public 4.7 `sm100.mma` atom exposes no A +//! collector field (its fields are accumulate, negate A/B and disable lanes). +//! The collector-aware MMA instruction therefore uses an NVVM architecture leaf: +//! its descriptors come from CuTe layouts and its shape/collector sequence come +//! from the verified tiled-MMA plan, never raw Rust instruction descriptors. + +use dialect_cute::{ + attributes::{CuteComposedLayoutAttr, CuteTmaStorePipelineAttr}, + sm100_ops::*, +}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; +use pliron_mlir_export::{ + DropAttribute, MlirAttribute, MlirBlock, MlirFloatType, MlirLocation, MlirOperation, + MlirRegion, MlirResult, MlirType, MlirValueUse, OperationInput, OperationTranslation, + TranslationError, TranslationRegistry, TranslationSession, +}; + +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)?; + 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(()) +} + +#[derive(Clone, Copy)] +enum Kind { + Alloc, + Dealloc, + Mma, + Epilogue, + TmaLoad, + PipelineInit, + PipelineAcquire, + PipelineExpect, + PipelineWait, + PipelineRelease, + AccInit, + AccAcquire, + AccCommit, + AccWait, + AccRelease, + TmaStore, + StoreCommit, + StoreAcquire, + StoreTail, +} +enum Plan { + Tmem(CuteSm100TmemPlanAttr), + Mma(CuteSm100MmaPlanAttr), + Epilogue(CuteSm100EpiloguePlanAttr), + Copy(CuteSm100ClusterCopyPlanAttr), + Pipeline(CuteSm100PipelinePlanAttr), + Store(CuteTmaStorePipelineAttr), +} +struct Recipe { + kind: Kind, + 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) + } +} + +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(()) +} + +fn instruction_descriptor(m: u32, n: u32) -> u32 { + (1 << 4) | ((n >> 3) << 17) | ((m >> 4) << 24) +} + +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.outer(); + let mask = e.constant(32, (((1u32 << swizzle.bits) - 1) << swizzle.base) as i128)?; + let shift = e.constant(32, swizzle.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(()) +} + +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(()) +} + +struct Emit<'a, 'b> { + session: &'a mut TranslationSession<'b>, + location: MlirLocation, + ops: Vec, +} +impl Emit<'_, '_> { + fn operation(&self, name: &str, operands: Vec) -> Result { + let mut op = MlirOperation::new(name)?; + op.operands = operands; + op.location = self.location.clone(); + Ok(op) + } + fn fresh(&mut self, ty: MlirType) -> (MlirResult, MlirValueUse) { + let id = self.session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) + } + fn emit(&mut self, name: &str, operands: Vec) -> Result<(), String> { + let op = self.operation(name, operands)?; + self.ops.push(op); + Ok(()) + } + 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) + } + fn cast( + &mut self, + name: &str, + operand: MlirValueUse, + ty: MlirType, + ) -> Result { + self.value(name, vec![operand], ty) + } + fn binary( + &mut self, + name: &str, + a: MlirValueUse, + b: MlirValueUse, + ) -> Result { + let ty = a.ty.clone(); + self.value(name, vec![a, b], ty) + } + 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) + } + 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) + } + } + 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}>" + ))?, + ) + } + 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) + } + 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), + } + } + 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.outer(); + let ptr = self.cute_ptr( + ptr, + "f16", + "smem", + 128, + Some(format!( + "S<{},{},{}>", + swizzle.bits, + swizzle.base + 1, + swizzle.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), + ) + } + 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(()) + } + fn region(&mut self, operations: Vec) -> MlirRegion { + MlirRegion { + blocks: vec![MlirBlock { + id: self.session.fresh_block(), + arguments: vec![], + operations, + }], + } + } + 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(()) + } +} +fn dialect(s: &str) -> Result { + MlirType::dialect(s) +} +fn llvm_ptr_type(space: u32) -> Result { + dialect(&if space == 0 { + "!llvm.ptr".into() + } else { + format!("!llvm.ptr<{space}>") + }) +} +fn property(op: &mut MlirOperation, key: &str, value: &str) -> Result<(), String> { + op.properties + .insert(key.into(), MlirAttribute::dialect(value)?); + Ok(()) +} +fn integer_property(op: &mut MlirOperation, key: &str, bits: u32, value: i128) { + op.properties.insert( + key.into(), + MlirAttribute::Integer { + value, + ty: MlirType::Integer(bits), + }, + ); +} +fn segments(op: &mut MlirOperation, values: &[i32]) { + op.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(values.to_vec()), + ); +} + +#[cfg(test)] +mod tests { + 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, + }; + + 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(), + ) + } + 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), + } + } + 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) + } + 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 = pliron_mlir_export::translate_module( + ctx, + module, + ®istry, + &profile.translation_config(), + ) + .unwrap(); + pliron_mlir_export::render_module(&target) + } + #[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}"); + } + } + #[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}"); + } + #[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)) + ); + } + + #[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/crates/cuda-oxide-mlir-export/src/lib.rs b/crates/cuda-oxide-mlir-export/src/lib.rs index fee9654554..3f05296b89 100644 --- a/crates/cuda-oxide-mlir-export/src/lib.rs +++ b/crates/cuda-oxide-mlir-export/src/lib.rs @@ -12,16 +12,22 @@ mod cute; mod cute_gemm; mod cute_gemv; +mod cute_sm100; mod mir_core; mod mir_memory; +mod nvvm_cluster; mod nvvm_sregs; +mod nvvm_tcgen05; mod packs; mod profile; pub use cute::register_cute_elementwise_pack; pub(crate) use cute_gemm::register_cute_gemm_pack; pub(crate) use cute_gemv::register_cute_cutlass_profile_pack; +pub(crate) use cute_sm100::register_cute_sm100_pack; pub use mir_core::register_mir_core_pack; +pub use nvvm_cluster::register_nvvm_cluster_pack; pub use nvvm_sregs::register_nvvm_sreg_pack; +pub use nvvm_tcgen05::register_nvvm_tcgen05_pack; pub use packs::register_builtin_pack; pub use profile::{CutlassFullCuteMlir22, ExportManifest, MlirConsumerProfile, ProfileError}; diff --git a/crates/cuda-oxide-mlir-export/src/mir_core.rs b/crates/cuda-oxide-mlir-export/src/mir_core.rs index ec5ec8ea3a..1516f98d23 100644 --- a/crates/cuda-oxide-mlir-export/src/mir_core.rs +++ b/crates/cuda-oxide-mlir-export/src/mir_core.rs @@ -63,8 +63,8 @@ pub fn register_mir_core_pack(registry: &mut TranslationRegistry) -> Result<(), registry.register_operation::(IntegerBinaryTranslation("arith.andi"))?; registry.register_operation::(IntegerBinaryTranslation("arith.ori"))?; registry.register_operation::(IntegerBinaryTranslation("arith.xori"))?; - registry.register_operation::(IntegerBinaryTranslation("arith.shli"))?; - registry.register_operation::(ShiftRightTranslation)?; + registry.register_operation::(ShiftTranslation { right: false })?; + registry.register_operation::(ShiftTranslation { right: true })?; registry.register_operation::(NegTranslation)?; registry.register_operation::(NotTranslation)?; @@ -136,6 +136,7 @@ impl OperationTranslation for FunctionTranslation { .attributes .insert("gpu.kernel".into(), MlirAttribute::Unit); move_exact_block_contract(&mut target)?; + move_cluster_contract(&mut target)?; flatten_kernel_slice_abi(ctx, source, &mut target, session)?; } other => { @@ -227,6 +228,25 @@ fn move_exact_block_contract(target: &mut MlirOperation) -> Result<(), String> { } } +/// 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()), + } +} + fn take_positive_i32_attribute( target: &mut MlirOperation, name: &str, @@ -531,23 +551,85 @@ impl OperationTranslation for IntegerBinaryTranslation { } } -struct ShiftRightTranslation; +struct ShiftTranslation { + right: bool, +} -impl OperationTranslation for ShiftRightTranslation { +impl OperationTranslation for ShiftTranslation { fn translate( &self, ctx: &Context, source: Ptr, - input: OperationInput, - _session: &mut TranslationSession<'_>, + mut input: OperationInput, + session: &mut TranslationSession<'_>, ) -> Result, String> { let integer = source_integer(ctx, source, 0)?; - let name = if integer.signedness() == Signedness::Signed { + let name = if !self.right { + "arith.shli" + } else if integer.signedness() == Signedness::Signed { "arith.shrsi" } else { "arith.shrui" }; - Ok(vec![renamed(name, input)?]) + 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) } } @@ -934,6 +1016,36 @@ fn move_property( mod tests { use std::num::NonZero; + #[test] + fn cluster_contract_rejects_partial_and_nonpositive_dimensions() { + use pliron_mlir_export::{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, @@ -1019,6 +1131,64 @@ mod tests { IntegerType::get(ctx, width, signedness) } + #[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}"); + } + } + } + #[test] fn scalar_arithmetic_keeps_signedness_in_the_operation_names() { let mut ctx = context(); diff --git a/crates/cuda-oxide-mlir-export/src/mir_memory.rs b/crates/cuda-oxide-mlir-export/src/mir_memory.rs index 57add5c211..e61625e966 100644 --- a/crates/cuda-oxide-mlir-export/src/mir_memory.rs +++ b/crates/cuda-oxide-mlir-export/src/mir_memory.rs @@ -10,11 +10,11 @@ //! field order. Operations then use LLVM's aggregate and memory operations. use dialect_mir::{ - attributes::FieldIndexAttr, + attributes::{CompilerResultBundleAttr, FieldIndexAttr}, ops::{ - MirAllocaOp, MirArrayElementAddrOp, MirAssertOp, MirCallOp, MirConstructStructOp, - MirExtractArrayElementOp, MirExtractFieldOp, MirFieldAddrOp, MirLoadOp, MirPtrOffsetOp, - MirStorageDeadOp, MirStorageLiveOp, MirStoreOp, + MirAllocaOp, MirArrayElementAddrOp, MirAssertOp, MirCallOp, MirConstructArrayOp, + MirConstructStructOp, MirExtractArrayElementOp, MirExtractFieldOp, MirFieldAddrOp, + MirLoadOp, MirPtrOffsetOp, MirStorageDeadOp, MirStorageLiveOp, MirStoreOp, }, types::{ MirArrayType, MirDisjointSliceType, MirFP16Type, MirPtrType, MirSliceType, MirStructType, @@ -24,6 +24,7 @@ use dialect_mir::{ 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}, @@ -46,6 +47,9 @@ pub(crate) fn register_mir_memory_pack( 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::(LoadTranslation)?; @@ -53,6 +57,7 @@ pub(crate) fn register_mir_memory_pack( registry.register_operation::(PointerOffsetTranslation)?; registry.register_operation::(ExtractFieldTranslation)?; registry.register_operation::(ConstructStructTranslation)?; + registry.register_operation::(ConstructArrayTranslation)?; registry.register_operation::(FieldAddressTranslation)?; registry.register_operation::(ArrayElementAddressTranslation)?; registry.register_operation::(ExtractArrayElementTranslation)?; @@ -366,6 +371,63 @@ impl OperationTranslation for ExtractFieldTranslation { struct ConstructStructTranslation; +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, @@ -1139,10 +1201,11 @@ mod tests { profile::render_mapping_module_without_cutlass_envelope, }; use dialect_mir::{ - attributes::FieldIndexAttr, + attributes::{CompilerResultBundleAttr, FieldIndexAttr}, ops::{ - MirAddOp, MirAllocaOp, MirArrayElementAddrOp, MirExtractArrayElementOp, - MirExtractFieldOp, MirFieldAddrOp, MirFuncOp, MirLoadOp, MirReturnOp, MirStoreOp, + MirAddOp, MirAllocaOp, MirArrayElementAddrOp, MirConstructArrayOp, + MirExtractArrayElementOp, MirExtractFieldOp, MirFieldAddrOp, MirFuncOp, MirLoadOp, + MirReturnOp, MirStoreOp, }, types::{MirArrayType, MirPtrType, MirStructType}, }; @@ -1178,6 +1241,58 @@ mod tests { IntegerType::get(ctx, width, Signedness::Unsigned) } + #[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(&mut 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}"); + } + } + #[test] fn aggregate_slots_and_dynamic_array_access_map_to_llvm() { let mut ctx = Context::new(); diff --git a/crates/cuda-oxide-mlir-export/src/nvvm_cluster.rs b/crates/cuda-oxide-mlir-export/src/nvvm_cluster.rs new file mode 100644 index 0000000000..231d6f29dd --- /dev/null +++ b/crates/cuda-oxide-mlir-export/src/nvvm_cluster.rs @@ -0,0 +1,1029 @@ +/* + * 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 dialect_mir::ops::MirConstantOp; +use dialect_nvvm::ops::*; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + op::Op, + operation::Operation, +}; +use pliron_mlir_export::{ + DropAttribute, MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, + OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, +}; + +/// Register the native primitives used by persistent two-CTA TMA kernels. +pub fn register_nvvm_cluster_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + 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!( + 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!(BarrierCtaSyncCountOp, "v1:i0883", CountedBarrier); + register!(MbarrierInitSharedOp, "v1:i0097", BarrierInit); + 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 } + ); + register!( + MbarrierTryWaitParityClusterOp, + "v1:i0311", + WaitParity { cluster: true } + ); + register!(MbarrierArriveClusterOp, "v1:i0308", RemoteArrive); + 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(()) +} + +enum Kind { + Sreg(&'static str), + NoOperands(&'static str), + AsyncFence, + CountedBarrier, + BarrierInit, + BarrierArrive, + ExpectTx { cluster: bool }, + WaitParity { cluster: bool }, + RemoteArrive, + MapShared, + Elect, + Prefetch, + TmaLoadCg2, + TmaStore, + WaitGroup { read: bool }, +} + +struct Recipe { + marker: &'static str, + kind: Kind, + 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::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 => { + signature(&input, &[Int(32), Int(32)], &[])?; + operation("nvvm.barrier.cta.sync", vec![], input.operands, location)? + } + Kind::BarrierInit => { + signature(&input, &[Pointer(&[0, 3]), Int(32)], &[])?; + operation("nvvm.mbarrier.init", 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 } => { + 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", "#nvvm.mem_order")?; + } + op + } + Kind::RemoteArrive => { + 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")?; + 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) + } +} + +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]) + } +} + +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(()) +} + +#[derive(Clone, Copy, Debug)] +enum TypeSpec { + Int(u32), + Pointer(&'static [u32]), +} + +impl TypeSpec { + 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)), + } + } +} + +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(()) +} + +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) +} + +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()) +} + +fn pointer_type(space: u32) -> Result { + MlirType::dialect(if space == 0 { + "!llvm.ptr".into() + } else { + format!("!llvm.ptr<{space}>") + }) +} + +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) +} + +fn dialect_property(op: &mut MlirOperation, name: &str, value: &str) -> Result<(), String> { + op.properties + .insert(name.into(), MlirAttribute::dialect(value)?); + Ok(()) +} + +fn fresh(session: &mut TranslationSession<'_>, ty: MlirType) -> (MlirResult, MlirValueUse) { + let id = session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) +} + +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) +} + +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 { + 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; + + struct Fixture { + ctx: Context, + module: ModuleOp, + entry: Ptr, + types: Vec, + u32: TypedHandle, + u64: TypedHandle, + } + + impl Fixture { + 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, + } + } + + fn arg(&self, index: usize) -> Value { + self.entry.deref(&self.ctx).get_argument(index) + } + + 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 + } + + 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) + } + + 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()) + } + } + + #[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") + ); + } + + #[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\"")); + } + + #[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}"); + } + + #[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}" + ); + } + + #[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 operands" + } else { + "zero cache hint" + }), + "{error}" + ); + } + } + + #[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/crates/cuda-oxide-mlir-export/src/nvvm_tcgen05.rs b/crates/cuda-oxide-mlir-export/src/nvvm_tcgen05.rs new file mode 100644 index 0000000000..ce624d165f --- /dev/null +++ b/crates/cuda-oxide-mlir-export/src/nvvm_tcgen05.rs @@ -0,0 +1,598 @@ +/* + * 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::{ + CvtaGenericToSharedOffsetOp, Tcgen05AllocCg2Op, Tcgen05CommitMulticastCg2Op, + Tcgen05DeallocCg2Op, Tcgen05FenceAfterThreadSyncOp, Tcgen05FenceBeforeThreadSyncOp, + Tcgen05Ld32x32bX32RawOp, Tcgen05LoadWaitOp, Tcgen05MmaBBufferAttr, Tcgen05MmaBUsageAttr, + Tcgen05MmaCollectorAAttr, Tcgen05MmaCtaGroupAttr, Tcgen05MmaFormAttr, Tcgen05MmaKindAttr, + Tcgen05MmaOp, Tcgen05RelinquishAllocPermitCg2Op, +}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; +use pliron_mlir_export::{ + DropAttribute, MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, + OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, +}; + +/// 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)?; + registry.register_operation::(LoadTranslation)?; + registry.register_operation::(PrimitiveTranslation(Primitive::Alloc))?; + registry.register_operation::(PrimitiveTranslation(Primitive::Dealloc))?; + registry.register_operation::(PrimitiveTranslation( + Primitive::Commit, + ))?; + registry.register_operation::(PrimitiveTranslation( + Primitive::Relinquish, + ))?; + registry.register_operation::(PrimitiveTranslation(Primitive::LoadWait))?; + registry.register_operation::(PrimitiveTranslation( + Primitive::FenceBefore, + ))?; + registry.register_operation::(PrimitiveTranslation( + Primitive::FenceAfter, + ))?; + Ok(()) +} + +struct MmaTranslation; + +struct SharedOffsetTranslation; + +impl OperationTranslation for SharedOffsetTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + 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 = 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())?; + if mma.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() != Some(&Tcgen05MmaFormAttr::Shared) + || mma.get_attr_nvvm_tcgen05_mma_kind(ctx).as_deref() != Some(&Tcgen05MmaKindAttr::F16) + || mma.get_attr_nvvm_tcgen05_mma_cta_group(ctx).as_deref() + != Some(&Tcgen05MmaCtaGroupAttr::Cg2) + { + return Err( + "CUTLASS TCGen05 mapping supports shared-A FP16 CTA-group-2 MMA only".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, "v1:i0763", 5, 0)?; + let (cast, tmem) = tmem_pointer(session, input.operands[0].clone(), &input.location)?; + input.operands[0] = tmem; + 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()?); + 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]), + ); + Ok(vec![cast, target]) + } +} + +#[derive(Clone, Copy)] +enum Primitive { + Alloc, + Dealloc, + Commit, + Relinquish, + LoadWait, + FenceBefore, + FenceAfter, +} + +struct PrimitiveTranslation(Primitive); + +impl OperationTranslation for PrimitiveTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let (name, marker, operands) = match self.0 { + Primitive::Alloc => { + Tcgen05AllocCg2Op::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + ("nvvm.tcgen05.alloc", "v1:i0359", 2) + } + Primitive::Dealloc => { + Tcgen05DeallocCg2Op::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + ("nvvm.tcgen05.dealloc", "v1:i0360", 2) + } + Primitive::Commit => { + Tcgen05CommitMulticastCg2Op::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + ("nvvm.tcgen05.commit", "v1:i0365", 2) + } + Primitive::Relinquish => { + Tcgen05RelinquishAllocPermitCg2Op::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + ("nvvm.tcgen05.relinquish_alloc_permit", "v1:i0361", 0) + } + Primitive::LoadWait => { + Tcgen05LoadWaitOp::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + ("nvvm.tcgen05.wait", "v1:i0357", 0) + } + Primitive::FenceBefore => { + Tcgen05FenceBeforeThreadSyncOp::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + ("nvvm.tcgen05.fence", "v1:i0346", 0) + } + Primitive::FenceAfter => { + Tcgen05FenceAfterThreadSyncOp::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + ("nvvm.tcgen05.fence", "v1:i0347", 0) + } + }; + require_input(&mut input, marker, operands, 0)?; + let mut operations = vec![]; + if matches!(self.0, Primitive::Dealloc) { + let (cast, pointer) = + tmem_pointer(session, input.operands[0].clone(), &input.location)?; + operations.push(cast); + input.operands[0] = pointer; + } + // Native SMEM pointer arguments can be generic after MIR casts. + // Preserve the address while making the required shared space explicit. + if matches!(self.0, Primitive::Alloc | Primitive::Commit) { + input.operands[0] = shared_pointer( + session, + &mut operations, + input.operands[0].clone(), + &input.location, + )?; + } + let mut target = operation(name, input.operands, vec![], &input.location)?; + let kind = match self.0 { + Primitive::LoadWait => Some("#nvvm.tcgen05_wait"), + Primitive::FenceBefore => Some("#nvvm.tcgen05_fence"), + Primitive::FenceAfter => Some("#nvvm.tcgen05_fence"), + _ => None, + }; + if let Some(kind) = kind { + target + .properties + .insert("kind".into(), MlirAttribute::dialect(kind)?); + } else { + target.properties.insert("group".into(), cta_group()?); + } + operations.push(target); + Ok(operations) + } +} + +struct LoadTranslation; + +impl OperationTranslation for LoadTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + Tcgen05Ld32x32bX32RawOp::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + require_input(&mut input, "v1:i0664", 1, 32)?; + let (cast, tmem) = tmem_pointer(session, input.operands[0].clone(), &input.location)?; + let mut operations = vec![cast]; + let (vector_result, vector) = fresh( + session, + MlirType::Vector { + shape: vec![32], + element: Box::new(MlirType::Integer(32)), + }, + ); + let mut load = operation( + "nvvm.tcgen05.ld", + vec![tmem], + vec![vector_result], + &input.location, + )?; + load.properties.insert( + "shape".into(), + MlirAttribute::dialect("#nvvm.tcgen05_ldst_shape")?, + ); + operations.push(load); + // `raw` is asynchronous: retain the caller's explicit tcgen05.wait. + // Do not insert a wait here and change the source scheduling policy. + 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); + } + Ok(operations) + } +} + +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(()) +} + +fn cta_group() -> Result { + MlirAttribute::dialect("#nvvm.cta_group") +} + +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, + )) +} + +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) +} + +fn fresh(session: &mut TranslationSession<'_>, ty: MlirType) -> (MlirResult, MlirValueUse) { + let id = session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) +} + +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 { + use super::*; + 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, + }; + use pliron_mlir_export::render_module; + + 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 + } + + 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(&mut ctx, Tcgen05MmaFormAttr::Shared); + mma.set_attr_nvvm_tcgen05_mma_kind(&mut ctx, Tcgen05MmaKindAttr::F16); + mma.set_attr_nvvm_tcgen05_mma_cta_group(&mut ctx, Tcgen05MmaCtaGroupAttr::Cg2); + mma.set_attr_nvvm_tcgen05_mma_collector_a(&mut 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()) + } + + 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())?; + Ok(render_module(&target)) + } + + #[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")); + } + + #[test] + fn unsupported_mma_forms_fail_instead_of_losing_their_selectors() { + let (mut ctx, module, first) = fixture(); + let mma = Tcgen05MmaOp::new(first); + mma.set_attr_nvvm_tcgen05_mma_kind(&mut ctx, Tcgen05MmaKindAttr::Tf32); + let error = translate(&ctx, &module).unwrap_err(); + assert!( + error.contains("shared-A FP16 CTA-group-2 MMA only"), + "{error}" + ); + mma.set_attr_nvvm_tcgen05_mma_kind(&mut ctx, Tcgen05MmaKindAttr::F16); + mma.set_attr_nvvm_tcgen05_mma_cta_group(&mut ctx, Tcgen05MmaCtaGroupAttr::Cg1); + let error = translate(&ctx, &module).unwrap_err(); + assert!( + error.contains("shared-A FP16 CTA-group-2 MMA only"), + "{error}" + ); + } + + #[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/crates/cuda-oxide-mlir-export/src/profile.rs b/crates/cuda-oxide-mlir-export/src/profile.rs index f8faf81150..f45652fb4d 100644 --- a/crates/cuda-oxide-mlir-export/src/profile.rs +++ b/crates/cuda-oxide-mlir-export/src/profile.rs @@ -14,7 +14,8 @@ use thiserror::Error; use crate::{ register_builtin_pack, register_cute_cutlass_profile_pack, register_cute_gemm_pack, - register_mir_core_pack, register_nvvm_sreg_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. @@ -123,8 +124,11 @@ impl MlirConsumerProfile for CutlassFullCuteMlir22 { 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)?; register_cute_cutlass_profile_pack(&mut registry)?; register_cute_gemm_pack(&mut registry)?; + register_cute_sm100_pack(&mut registry)?; registry.seal(); Ok(registry) } diff --git a/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs b/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs index 8494a7d568..416dc9b42e 100644 --- a/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs +++ b/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs @@ -294,7 +294,14 @@ fn production_kernel_gets_one_device_module_and_the_six_argument_launch_abi() { Identifier::try_from("gpu_kernel").unwrap(), StringAttr::new("true".into()), ); - for (name, value) in [("reqntid_x", 256), ("reqntid_y", 1), ("reqntid_z", 1)] { + for (name, value) in [ + ("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)), @@ -436,4 +443,9 @@ fn production_kernel_gets_one_device_module_and_the_six_argument_launch_abi() { text.contains("nvvm.reqntid = array"), "{text}" ); + assert!( + text.contains("nvvm.cluster_dim = array"), + "{text}" + ); + assert!(!text.contains("cluster_dim_x"), "{text}"); } diff --git a/crates/mir-importer/src/translator/cute/layout.rs b/crates/mir-importer/src/translator/cute/layout.rs index b808c9819d..25f354fce5 100644 --- a/crates/mir-importer/src/translator/cute/layout.rs +++ b/crates/mir-importer/src/translator/cute/layout.rs @@ -24,7 +24,11 @@ use rustc_public::mir; /// 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. -fn const_u64(c: &rustc_public::ty::TyConst, what: &str, loc: &Location) -> TranslationResult { +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() { diff --git a/crates/mir-importer/src/translator/cute/mod.rs b/crates/mir-importer/src/translator/cute/mod.rs index 496322274f..be3908054a 100644 --- a/crates/mir-importer/src/translator/cute/mod.rs +++ b/crates/mir-importer/src/translator/cute/mod.rs @@ -39,6 +39,7 @@ pub(crate) mod pipeline; pub(crate) mod pipeline_emit; pub(crate) mod scheduler; pub(crate) mod scheduler_emit; +pub(crate) mod sm100_emit; pub(crate) mod smem_mma; pub(crate) mod smem_mma_emit; pub(crate) mod static_config; @@ -60,6 +61,7 @@ use rustc_public::ty::FnDef; /// Recognized cute-rs entry points. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum CuteFn { + Sm100(sm100_emit::Sm100Fn), MakeTensorRead, MakeTensorWrite, ZippedDivideRead, @@ -141,6 +143,108 @@ pub(crate) struct CuteFnSpec { /// the op set stabilizes this graduates to catalog generation (the /// cuda-intrinsics-gen model) and these entries become generated. pub(crate) const CUTE_FNS: &[CuteFnSpec] = &[ + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_tma_store", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::TmaStore), + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_store_commit", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::StoreCommit), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_store_acquire", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::StoreAcquire), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_store_tail", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::StoreTail), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_tmem_alloc", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::TmemAlloc), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_tmem_dealloc", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::TmemDealloc), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_tiled_mma", + func: CuteFn::Sm100(sm100_emit::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_emit::Sm100Fn::TmemEpilogue), + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_cluster_tma_load", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::ClusterTmaLoad), + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_pipeline_init", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::PipelineInit), + schema: &[SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_pipeline_acquire", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::PipelineAcquire), + schema: &[SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_pipeline_expect", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::PipelineExpect), + schema: &[SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_pipeline_wait", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::PipelineWait), + schema: &[SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_pipeline_release", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::PipelineRelease), + schema: &[SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_accumulator_init", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::AccumulatorInit), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_accumulator_acquire", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::AccumulatorAcquire), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_accumulator_commit", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::AccumulatorCommit), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_accumulator_wait", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::AccumulatorWait), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_accumulator_release", + func: CuteFn::Sm100(sm100_emit::Sm100Fn::AccumulatorRelease), + schema: &[SubstKind::Const], + }, CuteFnSpec { canonical: "cute_rs::tensor::make_tensor_read", func: CuteFn::MakeTensorRead, @@ -471,6 +575,7 @@ const CUTE_SOURCE_ABSTRACTION_MODULES: &[&str] = &[ "numeric", "pipeline", "scheduler", + "sm100", "tensor", "tiled_copy", ]; @@ -551,6 +656,7 @@ mod canonical_path_tests { "numeric", "pipeline", "scheduler", + "sm100", "tensor", "tiled_copy", ] @@ -1471,6 +1577,9 @@ pub(crate) fn try_translate_cute_call( } Some(match cute_fn { + CuteFn::Sm100(kind) => sm100_emit::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, diff --git a/crates/mir-importer/src/translator/cute/sm100_emit.rs b/crates/mir-importer/src/translator/cute/sm100_emit.rs new file mode 100644 index 0000000000..1045902438 --- /dev/null +++ b/crates/mir-importer/src/translator/cute/sm100_emit.rs @@ -0,0 +1,358 @@ +/* + * 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::{Located, 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}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Sm100Fn { + TmemAlloc, + TmemDealloc, + TiledMma, + TmemEpilogue, + ClusterTmaLoad, + TmaStore, + StoreCommit, + StoreAcquire, + StoreTail, + PipelineInit, + PipelineAcquire, + PipelineExpect, + PipelineWait, + PipelineRelease, + AccumulatorInit, + AccumulatorAcquire, + AccumulatorCommit, + AccumulatorWait, + AccumulatorRelease, +} + +fn invalid(message: impl core::fmt::Display) -> pliron::result::Error { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid SM100 CuTe boundary: {message}" + ))) +} + +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) +} + +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")) +} + +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) +} + +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, + }) +} + +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, + }) +} + +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], + } +} + +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, + } +} + +#[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/cuteir/README.md b/cuteir/README.md index 038f41a649..326ab48302 100644 --- a/cuteir/README.md +++ b/cuteir/README.md @@ -57,6 +57,7 @@ form a complete semantic story. It does not choose or execute a backend. | [`elementwise_cute`](examples/elementwise_cute) | Adds two vectors through per-thread Tensor tiles. | | [`nvfp4_gemv_cute`](examples/nvfp4_gemv_cute) | Multiplies a packed FP4 matrix by a packed FP4 vector. | | [`blockscale_gemm_cute`](examples/blockscale_gemm_cute) | Multiplies two block-scaled FP4 matrices with the Tensor Memory Accelerator (TMA) and tensor cores. | +| [`fp16_gemm_256x352_cute`](examples/fp16_gemm_256x352_cute) | SM100 two-CTA GEMM with typed shared tiles, TMEM operations, collector reuse, and cluster TMA pipelines through the CUTLASS translation backend. | The examples are deliberately ordered from smallest to largest: @@ -84,8 +85,10 @@ The elementwise example uses ordinary `f32` and `f16` operations. The two FP4 examples use Blackwell instructions and should be built for `sm_120a`; their own READMEs give the exact device and shape requirements. -Each program checks every GPU output bit pattern against a deterministic host -result before reporting success. Backend-specific artifact and code-shape +The first three programs check every GPU output bit pattern against a deterministic host +result before reporting success. The FP16 example checks every output +against a numerical oracle; its README includes the SM100 build and validation commands. +Backend-specific artifact and code-shape checks live with the backend implementation rather than in this shared layer. On the translation branch, the [official CUTLASS backend](docs/translation.md) diff --git a/cuteir/cute-rs/src/lib.rs b/cuteir/cute-rs/src/lib.rs index 067fe522e4..7f00ea4035 100644 --- a/cuteir/cute-rs/src/lib.rs +++ b/cuteir/cute-rs/src/lib.rs @@ -66,6 +66,7 @@ pub mod mma; pub mod numeric; pub mod pipeline; pub mod scheduler; +pub mod sm100; pub mod tensor; pub mod tile; pub mod tiled_copy; @@ -93,6 +94,10 @@ pub use mma::{AccC, FragA, FragB, load_matrix_a, load_matrix_b}; pub use numeric::{E2M1, Mxf4E2M1, PackedE2M1x2, UE8M0, UE8M0x2, UE8M0x4}; pub use pipeline::{Consumer, PipelineState, Producer, TmaLoadPipeline, TmaStorePipeline}; pub use scheduler::{StaticPersistentTileScheduler, WorkTile}; +pub use sm100::{ + OperandA, OperandB, Sm100AccumulatorPipeline, Sm100ClusterTmaCopy, Sm100SharedTile, + Sm100TiledMma, Sm100TmaMmaPipeline, Sm100TmaStorePipeline, Sm100Tmem, Sm100TmemEpilogue, +}; pub use tensor::{Contiguous1D, RegisterTile, Tensor, TensorElement, TensorMut, Tile1D, Zipped1D}; pub use tile::{assume_div, load_tile, store_tile}; pub use tiled_copy::{GlobalCopyTensor, SharedTensor, TiledCopy}; diff --git a/cuteir/cute-rs/src/sm100.rs b/cuteir/cute-rs/src/sm100.rs new file mode 100644 index 0000000000..b21d26a19e --- /dev/null +++ b/cuteir/cute-rs/src/sm100.rs @@ -0,0 +1,491 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SM100 two-CTA tensor-memory GEMM building blocks. +//! +//! Shared layouts, the two N partitions, and the A collector lifetime are +//! compiler-visible semantics. The scalar compiler boundary below is lowered +//! to dialect-cute and then to CUTLASS; it has no raw CUDA implementation. + +use crate::TmaDesc; +use crate::markers::ReifySmem2D; +use core::marker::PhantomData; +use cuda_device::barrier::Barrier; + +/// A shared tile used as the left input of a matrix product. +pub enum OperandA {} +/// A shared tile used as the right input of a matrix product. +pub enum OperandB {} + +/// A non-owning FP16 shared tile with a layout and an operand role. +pub struct Sm100SharedTile { + base: *mut f16, + 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, + } + } +} + +/// 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 { + 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 allocate(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 deallocate(address: u32) { + let () = Self::VALID; + unsafe { __compiler::sm100_tmem_dealloc::(address) } + } +} + +/// Two-CTA FP16 tiled MMA with two N partitions sharing the A collector. +/// +/// Each K=16 step fills A with the first N partition and consumes it for +/// the last time with the second. The collector lifetime never crosses a K +/// step. The complete output is `256 x (N0 + N1)` with FP32 accumulation. +pub struct Sm100TiledMma( + PhantomData<(AL, B0L, B1L)>, +); +impl< + AL: ReifySmem2D, + B0L: ReifySmem2D, + B1L: ReifySmem2D, + const N0: u32, + const N1: u32, + const K: u32, +> Sm100TiledMma +{ + const VALID: () = { + assert!(N0 >= 16 && N0 <= 256 && N0 % 16 == 0); + assert!(N1 >= 16 && N1 <= 256 && N1 % 16 == 0); + assert!(N0 + N1 <= 512 && K == 64); + assert!(AL::ROWS == 128 && AL::COLS == K as i64); + assert!(B0L::ROWS == (N0 / 2) as i64 && B0L::COLS == K as i64); + assert!(B1L::ROWS == (N1 / 2) as i64 && B1L::COLS == K as i64); + }; + + /// Define the compile-time tile and collector contract. + #[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() + } +} + +/// A cluster TMA copy whose descriptor and destination have the same layout. +pub struct Sm100ClusterTmaCopy(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 + /// its issuing lane internally 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 TmaDesc, + 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) + } + } +} + +/// Drain one 128x32 FP32 accumulator slice to a laid-out FP16 shared tile. +pub struct Sm100TmemEpilogue(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 TmaDesc, + src: *const f16, + row: u32, + column: u32, + ) { + unsafe { __compiler::sm100_tma_store::(desc, src, row, column) } + } +} + +/// Two-CTA TMA-to-MMA stage ring, with one completion arrival per stage. +/// +/// `TX_BYTES` includes every copy from both CTAs. The elected leader produces +/// the expectation, and asynchronous MMA completion releases both CTAs. +pub struct Sm100TmaMmaPipeline { + full: *mut Barrier, + empty: *mut Barrier, +} +impl Sm100TmaMmaPipeline { + 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(&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 full_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 Sm100AccumulatorPipeline { + empty: *mut Barrier, + full: *mut Barrier, +} +impl Sm100AccumulatorPipeline { + 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. +#[derive(Clone, Copy, Default)] +pub struct Sm100TmaStorePipeline; +impl Sm100TmaStorePipeline { + 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::() } + } +} + +/// Stable semantic compiler boundaries; never execute on the host. +#[doc(hidden)] +#[allow(missing_docs)] +pub mod __compiler { + use super::*; + #[inline(never)] + pub unsafe fn sm100_tmem_alloc(token: *mut u32) { + let _ = token; + unreachable!("SM100 CuTe marker requires device compilation") + } + #[inline(never)] + pub unsafe fn sm100_tmem_dealloc(tmem: u32) { + let _ = tmem; + unreachable!("SM100 CuTe marker requires device compilation") + } + #[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") + } + #[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") + } + // Collective boundary: every lane of one producer warp participates. + // The native CuTe load owns election and any required warp synchronization. + #[inline(never)] + pub unsafe fn sm100_cluster_tma_load( + desc: *const TmaDesc, + 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") + } + + #[inline(never)] + pub unsafe fn sm100_tma_store( + desc: *const TmaDesc, + src: *const f16, + row: u32, + col: u32, + ) { + let _ = (desc, src, row, col); + unreachable!("SM100 CuTe marker requires device compilation") + } + macro_rules! store_marker { + ($name:ident) => { + #[inline(never)] + pub unsafe fn $name() { + unreachable!("SM100 CuTe marker requires device compilation") + } + }; + } + store_marker!(sm100_store_commit); + store_marker!(sm100_store_acquire); + store_marker!(sm100_store_tail); + macro_rules! pipeline_marker { + ($name:ident ($($arg:ident:$ty:ty),*)) => { + #[inline(never)] + pub unsafe fn $name($($arg:$ty),*) { + let _ = ($($arg),*); unreachable!("SM100 CuTe marker requires device compilation") + } + }; + } + pipeline_marker!(sm100_pipeline_init(full: *mut Barrier, empty: *mut Barrier)); + pipeline_marker!(sm100_pipeline_acquire(empty: *mut Barrier, stage: u32, phase: u32)); + pipeline_marker!(sm100_pipeline_expect(full: *mut Barrier, stage: u32)); + pipeline_marker!(sm100_pipeline_wait(full: *mut Barrier, stage: u32, phase: u32)); + pipeline_marker!(sm100_pipeline_release(empty: *mut Barrier, stage: u32)); + macro_rules! accumulator_marker { + ($name:ident ($($arg:ident:$ty:ty),*)) => { + #[inline(never)] + pub unsafe fn $name($($arg:$ty),*) { + let _ = ($($arg),*); unreachable!("SM100 CuTe marker requires device compilation") + } + }; + } + accumulator_marker!(sm100_accumulator_init(empty: *mut Barrier, full: *mut Barrier)); + accumulator_marker!(sm100_accumulator_acquire(empty: *mut Barrier, phase: u32)); + accumulator_marker!(sm100_accumulator_commit(full: *mut Barrier)); + accumulator_marker!(sm100_accumulator_wait(full: *mut Barrier, phase: u32)); + accumulator_marker!(sm100_accumulator_release(empty: *mut Barrier)); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{Composed, RowMajor, Swizzle}; + type A = Composed, 0, RowMajor<128, 64>>; + type B0 = Composed, 0, RowMajor<96, 64>>; + type B1 = Composed, 0, RowMajor<80, 64>>; + #[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/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.rs b/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.rs new file mode 100644 index 0000000000..a084526bd7 --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.rs @@ -0,0 +1,9 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#![feature(f16)] +use cute_rs::{Composed, RowMajor, Sm100TiledMma, Swizzle}; +type A = Composed, 0, RowMajor<128, 64>>; +type B0 = Composed, 0, RowMajor<96, 64>>; +type B1 = Composed, 0, RowMajor<80, 64>>; +const _: Sm100TiledMma = Sm100TiledMma::new(); +fn main() {} diff --git a/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.stderr b/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.stderr new file mode 100644 index 0000000000..7f10421f13 --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.stderr @@ -0,0 +1,11 @@ +error[E0080]: evaluation panicked: assertion failed: B0L::ROWS == (N0 / 2) as i64 && B0L::COLS == K as i64 + --> src/sm100.rs + | + | assert!(B0L::ROWS == (N0 / 2) as i64 && B0L::COLS == K as i64); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::Sm100TiledMma::, 0, cute_rs::RowMajor<128, 64>>, cute_rs::Composed, 0, cute_rs::RowMajor<80, 64>>, cute_rs::Composed, 0, cute_rs::RowMajor<96, 64>>, 192, 160>::VALID` failed here + +note: erroneous constant encountered + --> src/sm100.rs + | + | let () = Self::VALID; + | ^^^^^^^^^^^ diff --git a/cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.rs b/cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.rs new file mode 100644 index 0000000000..f46deb54e0 --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.rs @@ -0,0 +1,16 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#![feature(f16)] +use cute_rs::{Composed, OperandB, RowMajor, Sm100SharedTile, Sm100TiledMma, Swizzle}; +type A = Composed, 0, RowMajor<128, 64>>; +type B0 = Composed, 0, RowMajor<96, 64>>; +type B1 = Composed, 0, RowMajor<80, 64>>; +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/cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.stderr b/cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.stderr new file mode 100644 index 0000000000..6256b3e68b --- /dev/null +++ b/cuteir/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:14:21 + | +14 | 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.rs + | + | pub unsafe fn gemm( + | ^^^^ diff --git a/cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.rs b/cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.rs new file mode 100644 index 0000000000..1ab0f9c25c --- /dev/null +++ b/cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.rs @@ -0,0 +1,15 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#![feature(f16)] +use cute_rs::{ + Composed, OperandA, RowMajor, Sm100ClusterTmaCopy, Sm100SharedTile, Swizzle, TmaDesc, +}; +type A = Composed, 0, RowMajor<128, 64>>; +type B = Composed, 0, RowMajor<96, 64>>; +fn main() { + let desc: *const TmaDesc = 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/cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.stderr b/cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.stderr new file mode 100644 index 0000000000..7f5a32f6ad --- /dev/null +++ b/cuteir/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:13:46 + | +13 | 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.rs + | + | pub unsafe fn copy( + | ^^^^ diff --git a/cuteir/cute-rs/tests/type_safety.rs b/cuteir/cute-rs/tests/type_safety.rs index 4e233bc99b..aeb1b1c552 100644 --- a/cuteir/cute-rs/tests/type_safety.rs +++ b/cuteir/cute-rs/tests/type_safety.rs @@ -15,4 +15,7 @@ fn tensor_and_block_scaled_types_fail_closed() { 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"); } diff --git a/cuteir/dialect-cute/src/lib.rs b/cuteir/dialect-cute/src/lib.rs index 07d1dd8857..55775320d2 100644 --- a/cuteir/dialect-cute/src/lib.rs +++ b/cuteir/dialect-cute/src/lib.rs @@ -16,8 +16,9 @@ //! //! 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, and the shared-tensor/tiled-MMA compute story. These semantic -//! operations are later consumed by the selected backend continuation. +//! stage copies, shared-tensor/tiled-MMA compute, and SM100 two-CTA tensor +//! memory with collector-aware MMA, cluster TMA, and asynchronous pipelines. +//! These semantic operations are consumed by the selected backend continuation. pub mod attributes; pub mod epilogue_ops; @@ -26,6 +27,7 @@ pub mod gemv_ops; pub mod ops; pub mod pipeline_ops; pub mod scheduler_ops; +pub mod sm100_ops; pub mod smem_mma_ops; pub mod tensor_ops; pub mod types; @@ -50,6 +52,8 @@ pub fn register(ctx: &mut Context) { &DialectName::try_new(CUTE_DIALECT_NAME).expect("valid dialect name"), ); + sm100_ops::register(ctx); + // The #[pliron_op]/#[pliron_attr] macros auto-register via Context::default; // explicit registration is the idempotent house convention. ops::CuteCopyOp::register(ctx); diff --git a/cuteir/dialect-cute/src/sm100_ops.rs b/cuteir/dialect-cute/src/sm100_ops.rs new file mode 100644 index 0000000000..185f0b7b10 --- /dev/null +++ b/cuteir/dialect-cute/src/sm100_ops.rs @@ -0,0 +1,913 @@ +/* + * 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. + +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::location::Located; +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}; + +/// 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 { + AFillLastUse, +} + +#[pliron_attr( + name = "cute.sm100_tmem_plan", + format = "`<` $columns `,` $cta_group `>`" +)] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub struct CuteSm100TmemPlanAttr { + pub columns: u32, + 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(()) + } +} + +#[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 { + pub m: u32, + pub n0: u32, + pub n1: u32, + pub k: u32, + pub tmem_columns: u32, + pub cta_group: u32, + pub collector: CuteSm100CollectorAttr, + pub a_layout: CuteComposedLayoutAttr, + pub b0_layout: CuteComposedLayoutAttr, + pub b1_layout: CuteComposedLayoutAttr, +} + +fn row_major(rows: u32, columns: u32, swizzle_bits: u32) -> ComposedLayout { + let inner: Layout = format!("({rows},{columns}):({columns},1)") + .parse() + .expect("fixed rank-two layout"); + ComposedLayout::new( + Swizzle::new(swizzle_bits, 3, 3), + 0, + inner, + OffsetUnit::Elements, + ) + .expect("fixed SM100 FP16 swizzle") +} + +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 % 16 != 0 + || self.n1 % 16 != 0 + || 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(()) + } +} + +#[pliron_attr( + name = "cute.sm100_cluster_copy_plan", + format = "`<` $layout `,` $cta_group `,` $leader_rank `>`" +)] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteSm100ClusterCopyPlanAttr { + pub layout: CuteComposedLayoutAttr, + pub cta_group: u32, + pub leader_rank: u32, +} +impl Verify for CuteSm100ClusterCopyPlanAttr { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + self.layout.verify(ctx)?; + let dimensions = self.layout.0.inner().shape.leaves(); + 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(()) + } +} + +#[pliron_attr( + name = "cute.sm100_epilogue_plan", + format = "`<` $layout `,` $rows `,` $columns `>`" +)] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteSm100EpiloguePlanAttr { + pub layout: CuteComposedLayoutAttr, + pub rows: u32, + 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(()) + } +} + +#[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 { + pub stages: u32, + pub transaction_bytes: u32, + pub consumer_arrivals: u32, + pub cta_group: u32, + pub leader_rank: u32, + 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 % 2 != 0 + { + 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 % 16 != 0 { + return verify_err_noloc!( + "SM100 input pipeline requires one MMA completion arrival and a transaction size divisible by 16" + ); + } + Ok(()) + } +} + +#[derive(Clone, Copy)] +enum Carrier { + U32, + Bool, + F32, + SharedU32, + SharedBarrier, + SharedF16, + ReadSharedF16, + Descriptor, +} +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, + } + } + } +} + +macro_rules! semantic_op { + ($ty:ident, $name:literal, $key:ident, $plan:ty, [$($carrier:ident),*], $check:expr) => { + #[pliron_op(name = $name, format, attributes = ($key: $plan))] + pub struct $ty; + impl $ty { + 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 + } + pub fn wrap(op: Ptr) -> Self { Self { op } } + 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() + } + 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!( + CuteSm100TmemAllocOp, + "cute.sm100_tmem_alloc", + sm100_tmem_alloc_plan, + CuteSm100TmemPlanAttr, + [SharedU32], + |_| true +); +semantic_op!( + CuteSm100TmemDeallocOp, + "cute.sm100_tmem_dealloc", + sm100_tmem_dealloc_plan, + CuteSm100TmemPlanAttr, + [U32], + |_| true +); +semantic_op!( + CuteSm100TiledMmaOp, + "cute.sm100_tiled_mma", + sm100_tiled_mma_plan, + CuteSm100MmaPlanAttr, + [U32, ReadSharedF16, ReadSharedF16, ReadSharedF16, Bool], + |_| true +); +semantic_op!( + CuteSm100TmemEpilogueOp, + "cute.sm100_tmem_epilogue", + sm100_tmem_epilogue_plan, + CuteSm100EpiloguePlanAttr, + [U32, SharedF16, U32, F32], + |_| true +); +semantic_op!( + CuteSm100ClusterTmaLoadOp, + "cute.sm100_cluster_tma_load", + sm100_cluster_tma_load_plan, + CuteSm100ClusterCopyPlanAttr, + [Descriptor, SharedF16, U32, U32, SharedBarrier, U32], + |_| true +); +semantic_op!( + CuteSm100PipelineInitOp, + "cute.sm100_pipeline_init", + sm100_pipeline_init_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, SharedBarrier], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes > 0 +); +semantic_op!( + CuteSm100PipelineAcquireOp, + "cute.sm100_pipeline_acquire", + sm100_pipeline_acquire_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, U32, U32], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes > 0 +); +semantic_op!( + CuteSm100PipelineExpectOp, + "cute.sm100_pipeline_expect", + sm100_pipeline_expect_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, U32], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes > 0 +); +semantic_op!( + CuteSm100PipelineWaitOp, + "cute.sm100_pipeline_wait", + sm100_pipeline_wait_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, U32, U32], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes > 0 +); +semantic_op!( + CuteSm100PipelineReleaseOp, + "cute.sm100_pipeline_release", + sm100_pipeline_release_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, U32], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes > 0 +); +semantic_op!( + CuteSm100AccumulatorInitOp, + "cute.sm100_accumulator_init", + sm100_accumulator_init_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, SharedBarrier], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes == 0 +); +semantic_op!( + CuteSm100AccumulatorAcquireOp, + "cute.sm100_accumulator_acquire", + sm100_accumulator_acquire_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, U32], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes == 0 +); +semantic_op!( + CuteSm100AccumulatorCommitOp, + "cute.sm100_accumulator_commit", + sm100_accumulator_commit_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes == 0 +); +semantic_op!( + CuteSm100AccumulatorWaitOp, + "cute.sm100_accumulator_wait", + sm100_accumulator_wait_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, U32], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes == 0 +); +semantic_op!( + CuteSm100AccumulatorReleaseOp, + "cute.sm100_accumulator_release", + sm100_accumulator_release_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes == 0 +); +semantic_op!( + CuteSm100TmaStoreOp, + "cute.sm100_tma_store", + sm100_tma_store_plan, + CuteSm100EpiloguePlanAttr, + [Descriptor, ReadSharedF16, U32, U32], + |_| true +); +semantic_op!( + CuteSm100StoreCommitOp, + "cute.sm100_store_commit", + sm100_store_commit_plan, + CuteTmaStorePipelineAttr, + [], + |_| true +); +semantic_op!( + CuteSm100StoreAcquireOp, + "cute.sm100_store_acquire", + sm100_store_acquire_plan, + CuteTmaStorePipelineAttr, + [], + |_| true +); +semantic_op!( + CuteSm100StoreTailOp, + "cute.sm100_store_tail", + sm100_store_tail_plan, + CuteTmaStorePipelineAttr, + [], + |_| true +); + +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 { + 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) { + 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); + 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 +} + +fn static_plan(ctx: &Context, op: Ptr) -> Option { + op.deref(ctx) + .attributes + .0 + .values() + .find_map(|attr| attr.downcast_ref::().cloned()) +} + +/// 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")?; + for needed in [ + CuteSm100TmemAllocOp::get_opid_static(), + CuteSm100TmemDeallocOp::get_opid_static(), + 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(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use dialect_mir::ops::MirUndefOp; + use pliron::basic_block::BasicBlock; + use pliron::builtin::ops::ModuleOp; + use pliron::builtin::types::Signedness; + use pliron::linked_list::ContainsLinkedList; + use pliron::r#type::TypeHandle; + + 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)), + } + } + 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, + } + } + 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) + } + 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(); + 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) + } + + #[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()); + } + #[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()); + } + #[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(); + } + #[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}" + ); + } + #[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") + ); + } + #[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") + ); + } + #[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/cuteir/dialect-cute/src/verify.rs b/cuteir/dialect-cute/src/verify.rs index d9c2f4eb55..f5e8162c0a 100644 --- a/cuteir/dialect-cute/src/verify.rs +++ b/cuteir/dialect-cute/src/verify.rs @@ -205,6 +205,7 @@ fn all_high_level_semantic_ids() -> Vec { ids.extend(tma_semantic_ids()); ids.extend(smem_mma_semantic_ids()); ids.extend(epilogue_semantic_ids()); + ids.extend(crate::sm100_ops::semantic_ids()); ids } @@ -216,7 +217,9 @@ fn verify_local_op(ctx: &Context, operation: Ptr) -> Result { }; } - let result = if opid == CuteTensorMakeOp::get_opid_static() { + let result = if let Some(result) = crate::sm100_ops::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)) @@ -3834,6 +3837,7 @@ pub fn verify_cute_semantics(ctx: &Context, module: Ptr) -> Result { verify_gemv_story(ctx, &scope)?; verify_tensor_story(ctx, &scope)?; verify_epilogue_story(ctx, &scope)?; + crate::sm100_ops::verify_story(ctx, &scope).map_err(invalid)?; } Ok(()) } diff --git a/cuteir/docs/assets/cutlass-translation-flow.svg b/cuteir/docs/assets/cutlass-translation-flow.svg new file mode 100644 index 0000000000..056d46a587 --- /dev/null +++ b/cuteir/docs/assets/cutlass-translation-flow.svg @@ -0,0 +1,47 @@ + + Rust CuTe compilation through CUTLASS + Read left to right across the top, then right to left across the bottom. cute-rs layouts, tiles, and pipelines become dialect-cute operations. Common preparation and whole-module verification precede CUTLASS MLIR mapping into CuTe, CuTeNVGPU, and reviewed architecture leaves. The pinned libCutlassCompiler 4.7 produces a validated cubin embedded in the host artifact for CUDA launch. + + + + + + + + + + + + + + + + + + + + + + + + cute-rs + dialect-cute + Prepare + verify + CUTLASS MLIR + libCutlassCompiler + Validated cubin + + + Rust layouts, tiles, + and pipelines + Semantic operations + Common preparation + Whole-module verification + CuTe / CuTeNVGPU + + architecture leaves + Pinned version 4.7 + Embedded host artifact + CUDA launch + + + diff --git a/cuteir/docs/translation.md b/cuteir/docs/translation.md index 5c37d830b4..c54ad920e7 100644 --- a/cuteir/docs/translation.md +++ b/cuteir/docs/translation.md @@ -1,77 +1,97 @@ # Official CUTLASS translation backend -The translation backend preserves the high-level `dialect-cute` program after -common preparation, verifies the complete semantic story, and maps it directly -to the MLIR dialects accepted by NVIDIA's official CUTLASS 4.7 compiler. +Translate Rust CuTe contracts into NVIDIA's pinned CUTLASS 4.7 MLIR profile: -```text -cute-rs kernel - │ - ▼ -high-level dialect-cute - │ - ├── common preparation - └── backend-neutral whole-module verification - │ - ▼ -CUTLASS full-CuTe MLIR profile - │ - ▼ -official libCutlassCompiler 4.7 - │ - ▼ -validated cubin → embedded host artifact → CUDA launch -``` - -The backend contains direct mapping packs for all three shared examples: - -- tensor tiling and copying for `elementwise_cute`; -- scaled tensor views and NVFP4 GEMV for `nvfp4_gemv_cute`; -- scheduler, work tiles, TMA load pipelines, shared-memory MMA, and epilogue - stores for `blockscale_gemm_cute`. +![CuTe compiler flow from Rust layouts through verified MLIR and CUTLASS to a CUDA cubin](assets/cutlass-translation-flow.svg) -It does not run a native CuTe expansion pass or lower through backend-specific -MIR/NVVM leaf intrinsics. In the GEMM epilogue, -`ReadyForTma` directly emits the async-shared proxy publication fence followed -by the counted CTA hand-off barrier. `Reusable` emits only the counted barrier. +- High-level CuTe survives preparation; reviewed mapping packs translate it + directly, without native CuTe expansion or prior MIR/NVVM leaf lowering. +- Ordinary non-CuTe kernels retain the default MIR/NVVM/LLVM-to-PTX backend. -## Install and select the compiler +## Install and build -Install the pinned compiler archive once: +From the repository root after normal CUDA Oxide setup: ```bash cargo oxide toolchain install cutlass -``` +export CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir -The installer verifies the official archive and `libCutlassCompiler.so` -digests. With no explicit compiler path, `cargo oxide` resolves this managed -installation and includes its library digest in the build fingerprint. +cargo oxide build elementwise_cute --arch sm_120a +cargo oxide build nvfp4_gemv_cute --arch sm_120a +cargo oxide build blockscale_gemm_cute --arch sm_120a +cargo oxide build fp16_gemm_256x352_cute --arch sm_100a +``` -Select the backend for a build: +- Installation verifies both archive and library digests. `cargo oxide` + resolves the managed compiler and fingerprints its library digest. +- An explicit absolute path overrides discovery; library loading still + enforces the pinned digest: ```bash CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ +CUDA_OXIDE_CUTLASS_COMPILER=/absolute/path/to/libCutlassCompiler.so \ cargo oxide build elementwise_cute --arch sm_120a +``` -CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ - cargo oxide build nvfp4_gemv_cute --arch sm_120a +## Supported mappings -CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ - cargo oxide build blockscale_gemm_cute --arch sm_120a -``` +| Example | Compiler-visible operations | +| --- | --- | +| [Elementwise](../examples/elementwise_cute) | Tensor tiling and copying | +| [NVFP4 GEMV](../examples/nvfp4_gemv_cute) | Scaled tensor views and GEMV | +| [Block-scaled GEMM](../examples/blockscale_gemm_cute) | Scheduler, work tiles, TMA pipelines, shared-memory MMA, epilogue stores | +| [SM100 FP16 GEMM](../examples/fp16_gemm_256x352_cute) | Typed shared tiles, TMEM, paired two-CTA MMA, cluster TMA, producer/consumer pipelines | -An explicit absolute library path overrides the managed installation: +Block-scaled epilogue hand-offs: `ReadyForTma` emits an async-shared proxy +publication fence, then a counted CTA barrier; `Reusable` emits only the barrier. -```bash -CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ -CUDA_OXIDE_CUTLASS_COMPILER=/absolute/path/to/libCutlassCompiler.so \ - cargo oxide build elementwise_cute --arch sm_120a +## SM100 contracts + +The [Rust APIs](../cute-rs/src/sm100.rs) become verified `cute.sm100_*` plans: + +| API | Contract | +| --- | --- | +| `Sm100SharedTile`, `TmaDesc` | Operand role and matching host/device layout | +| `Sm100TiledMma` | Two N partitions sharing A's collector lifetime | +| `Sm100TmaMmaPipeline`, `Sm100AccumulatorPipeline` | Buffer ownership and completion | +| `Sm100Tmem`, `Sm100TmemEpilogue` | TMEM allocation and FP32 → FP16 epilogue | + +- **Verified:** layout/partition compatibility, cluster transaction bytes, + barrier arrivals, and allocation/MMA/copy/release lifecycle presence. +- **Caller-owned unsafe preconditions:** dynamic phases, lane participation, + and pointer lifetimes. +- **Profile:** two CTAs; FP16 inputs, FP32 accumulation; M256/K64; + 128-byte input swizzle; 512 TMEM columns; 128×32 epilogue with 64-byte swizzle. + The example uses N192+160 and checks CPU FP32 accumulation rounded to FP16. + +### CuTe and the collector leaf + +- Native CuTe builds shared descriptors, TMEM operations, and cluster TMA. +- CUTLASS 4.7's native MMA atom lacks A-collector fill/last-use control. + The backend therefore emits an **NVVM TCGen05 MMA leaf** from the verified plan: + +```text +Each K16: first N partition → fill A collector + second partition → last use of A collector ``` -Use `CUDA_OXIDE_MLIR_OUTPUT=/absolute/path/module.mlir` for an optional textual -observation. MLIR dumps and generated cubins are build artifacts and must not be -committed. +- Translation owns descriptor bits, instruction selection, barrier arrivals, + completion multicast, and TMEM fences. + +### Collective TMA loads + +- `Sm100ClusterTmaCopy::copy`: all 32 producer lanes per CTA participate, + with identical arguments within each warp; native CuTe elects internally. +- An outer elected-lane branch can deadlock on stage reuse. Regression cases + cover more than five K stages and multiple persistent tiles. +- TMA **stores** have one issuing thread; the collective load rule does not apply. + +## Diagnostics and limits -The ordinary CUDA Oxide MIR/NVVM/LLVM-to-PTX backend remains the default for -non-CuTe kernels. Selecting the CUTLASS backend changes only the continuation -after shared preparation and semantic verification. +- `CUDA_OXIDE_DEVICE_BACKEND` accepts `native` or `cutlass-mlir`; unknown values + fail. `CUDA_OXIDE_BACKEND` selects the rustc backend library, not this path. +- Unsupported semantics/profiles, verification failures, or compiler failures + stop the build; translation does not fall back to another backend. +- To inspect prepared MLIR, add `CUDA_OXIDE_MLIR_OUTPUT="$PWD/module.mlir"` + to a CUTLASS build. MLIR dumps and generated cubins are build artifacts; + do not commit them. diff --git a/cuteir/examples/fp16_gemm_256x352_cute/.gitignore b/cuteir/examples/fp16_gemm_256x352_cute/.gitignore new file mode 100644 index 0000000000..5613b74e05 --- /dev/null +++ b/cuteir/examples/fp16_gemm_256x352_cute/.gitignore @@ -0,0 +1,13 @@ +# Local Python reference and comparison harness. +/reference/ +/compare.py + +# Local performance investigations and benchmark artifacts. +/benchmark-results/ +/benchmarks/ +/PERFORMANCE.md +/performance.md +/LARGE_WORKLOAD*.md +/large_workload*.md +*.nsys-rep +*.ncu-rep diff --git a/cuteir/examples/fp16_gemm_256x352_cute/Cargo.lock b/cuteir/examples/fp16_gemm_256x352_cute/Cargo.lock new file mode 100644 index 0000000000..e55d8136c7 --- /dev/null +++ b/cuteir/examples/fp16_gemm_256x352_cute/Cargo.lock @@ -0,0 +1,598 @@ +# 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.71.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn 2.0.119", +] + +[[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.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[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 = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.2.1" +dependencies = [ + "bindgen", +] + +[[package]] +name = "cuda-core" +version = "0.2.1" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-macros", + "half", + "oxide-artifacts", +] + +[[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", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn 2.0.119", +] + +[[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 = [ + "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.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 = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "fp16_gemm_256x352_cute" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "cuda-host", + "cute-rs", + "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.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 = "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.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[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 = "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 = "oxide-artifacts" +version = "0.2.1" +dependencies = [ + "object", +] + +[[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 = "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 = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[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.5", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +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.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12df2e0110f65b775f769bb17ef989067a1d931b2eb822bd4346631eeada89f9" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.5", +] + +[[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 = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "zerocopy" +version = "0.8.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d35102a9f36d089ccae9e4c6802bc118be4487b80aaffc0ab4e0cf5ce92d2873" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "146c01f5ab44258da43cf276c74a2763db2ff3969c9c652c3f2de07041d0b2bc" +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/cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml b/cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml new file mode 100644 index 0000000000..18a2bddb4a --- /dev/null +++ b/cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "fp16_gemm_256x352_cute" +version = "0.1.0" +edition = "2024" +license = "LicenseRef-NvidiaProprietary" +description = "SM100 two-CTA persistent FP16 GEMM with a 256x352 output tile" + +[workspace] + +[dependencies] +cute-rs = { path = "../../cute-rs", features = ["host"] } +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../crates/cuda-core" } +serde_json = "1" diff --git a/cuteir/examples/fp16_gemm_256x352_cute/README.md b/cuteir/examples/fp16_gemm_256x352_cute/README.md new file mode 100644 index 0000000000..dc2dc2441f --- /dev/null +++ b/cuteir/examples/fp16_gemm_256x352_cute/README.md @@ -0,0 +1,116 @@ +# FP16 GEMM with CuTe on SM100 + +A persistent, two-CTA GEMM expressed with Rust layouts, tiled copies/MMA, +tensor memory, and asynchronous pipelines. + +```text +C[M, N] = A[M, K] @ B[N, K].T + bias[M, 1] # bias is optional +``` + +- A, B, C, and bias are FP16; multiplication accumulates in FP32. +- `--has-bias` adds one value per output row before conversion to FP16. +- Layouts and pipeline contracts stay visible through the compiler: + +![CuTe compiler flow from Rust layouts through verified MLIR and CUTLASS to a CUDA cubin](../../docs/assets/cutlass-translation-flow.svg) + +## Tile and pipeline + +```text + one 256 × 352 output tile + 192 columns 160 columns + ┌─────────────────┬────────────────┐ +CTA 0: 128 rows│ │ │ + ├─────────────────┼────────────────┤ +CTA 1: 128 rows│ │ │ + └─────────────────┴────────────────┘ + two-CTA MMA pair; A collector reused + +A/B global ──cluster TMA──▶ shared [5 stages] ──MMA──▶ TMEM +C global ◀─────TMA────── shared [2 stages] ◀──FP16 + bias──┘ +``` + +- **Cooperative MMA:** 256×192 and 256×160 outputs share A for each K16 + instruction pair; each mainloop stage covers K64. +- **Warp roles:** each CTA has 192 threads: four epilogue warps, one MMA + warp, and one TMA warp. +- **Input pipeline:** five shared stages overlap cluster TMA and MMA. + Every lane of the producer warp participates in each native CuTe copy. +- **Accumulator and epilogue:** 512 TMEM columns, then eleven 128×32 slices + per CTA; bias/conversion and vector stores feed two shared C stages. +- **Persistent schedule:** clusters visit tiles in M-first order. The host + uses the smaller of tile count and resident capacity; `--clusters` overrides it. +- **Shared memory:** 211,072 bytes per CTA, aligned to 1,024 bytes. Host + tensor maps and device tiles use the same Rust layout types: + +| Layout | Shared tile (rows × columns) | TMA swizzle | +| --- | --- | --- | +| `ASmem` | 128 × 64 | 128 bytes | +| `B0Smem` | 96 × 64 | 128 bytes | +| `B1Smem` | 80 × 64 | 128 bytes | +| `CSmem` | 128 × 32 | 64 bytes | + +## Build and run + +Requirements: an SM100 GPU, a compatible CUDA driver, and the repository's +configured cuda-oxide toolchain. From the repository root: + +```bash +cargo oxide toolchain install cutlass +export CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir +cargo oxide build fp16_gemm_256x352_cute --arch sm_100a + +# Exactly one kernel launch; verify its output afterward. +cargo oxide run fp16_gemm_256x352_cute --arch sm_100a -- \ + --mnk 256,352,64 --timing-mode direct --warmup 0 --iters 1 + +# N/K tails, row bias, persistent tiles, and five-stage ring reuse. +cargo oxide run fp16_gemm_256x352_cute --arch sm_100a -- \ + --mnk 512,712,392 --has-bias --clusters 1 \ + --timing-mode direct --warmup 0 --iters 1 +``` + +- M must be a positive multiple of **256**; N and K, of **8**. + TMA handles partial N352 and K64 tiles. +- The runner checks every output against CPU FP32 accumulation rounded to + FP16 (`atol=0.1`, `rtol=1e-5`). Large shapes make this check expensive. +- Standalone timing defaults to **graph**. The commands above explicitly use + **direct**: one kernel per CUDA-event pair, with no graph capture. +- `--warmup 0 --iters 1` in direct mode launches once total and verifies that + timed output. Setup and verification are outside the event interval. +- `--json result.json` saves configuration, verification, and timing metadata; + `--help` lists all options. Graph mode uses `--graph-launches` launches per sample. + +## Rust API example + +These declarations are taken from [the kernel](src/kernel.rs). The layout +parameters bind shared operands, tensor maps, and the paired MMA dimensions: + +```rust +use cute_rs::{Composed, RowMajor, Sm100TiledMma, Sm100TmaMmaPipeline, Swizzle}; + +type ASmem = Composed, 0, RowMajor<128, 64>>; +type B0Smem = Composed, 0, RowMajor<96, 64>>; +type B1Smem = Composed, 0, RowMajor<80, 64>>; +type TiledMma = Sm100TiledMma; +type MainloopPipeline = Sm100TmaMmaPipeline<5, 77_824>; +``` + +- [SM100 APIs](../../cute-rs/src/sm100.rs): `Sm100SharedTile`, + `Sm100ClusterTmaCopy`, `Sm100TiledMma`, `Sm100TmemEpilogue`, and pipelines. +- [Host setup](src/main.rs): typed TMA descriptors, cluster launch, and validation. +- [Dialect operations](../../dialect-cute/src/sm100_ops.rs): layout, collector, + and pipeline verification before lowering. + +## Regression checks + +[validate.py](validate.py) uses only Python's standard library and runs six +cases covering full tiles, bias, tails, persistent scheduling, and pipeline +phase changes. It checks CPU results and graph replay output on Linux/SM100. + +```bash +python cuteir/examples/fp16_gemm_256x352_cute/validate.py --sanitizer +``` + +- Build first; use `--binary PATH` if Cargo has a custom target directory. +- `--sanitizer` adds Compute Sanitizer memcheck/TMA and synccheck checks; + omit it to run only the six cases. diff --git a/cuteir/examples/fp16_gemm_256x352_cute/src/kernel.rs b/cuteir/examples/fp16_gemm_256x352_cute/src/kernel.rs new file mode 100644 index 0000000000..1e0b3d7a86 --- /dev/null +++ b/cuteir/examples/fp16_gemm_256x352_cute/src/kernel.rs @@ -0,0 +1,297 @@ +// {$nv-internal-release file} +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: LicenseRef-NvidiaProprietary +// +// Rust port of NVIDIA DKG fp16_gemm_3_256x352.py (revision b3f067d4). +// NVIDIA CORPORATION, its +// affiliates and licensors retain all intellectual property and proprietary +// rights in and to this material and modifications thereto. Use, reproduction, +// disclosure or distribution requires an express NVIDIA license agreement. + +//! 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 cuda_device::{ + DynamicSharedArray, cluster_launch, cuda_module, kernel, launch_bounds, launch_contract, +}; +use cuda_device::{barrier::*, cluster, thread, tma::prefetch_tma_descriptor, warp}; +use cute_rs::{ + Composed, OperandA, OperandB, RowMajor, Sm100AccumulatorPipeline, Sm100ClusterTmaCopy, + Sm100SharedTile, Sm100TiledMma, Sm100TmaMmaPipeline, Sm100TmaStorePipeline, Sm100Tmem, + Sm100TmemEpilogue, Swizzle, TmaDesc, +}; + +// These element layouts also encode the host TMA descriptors. FP16 shifts +// the swizzle's low-bit position by one when lowering to byte addresses. +pub type ASmem = Composed, 0, RowMajor<128, 64>>; +pub type B0Smem = Composed, 0, RowMajor<96, 64>>; +pub type B1Smem = Composed, 0, RowMajor<80, 64>>; +pub type CSmem = Composed, 0, RowMajor<128, 32>>; + +pub const THREADS: u32 = 192; +const STAGES: usize = 5; +const A_BYTES: usize = 128 * 64 * 2; +const B_BYTES: usize = 176 * 64 * 2; +const B_OFFSET: usize = A_BYTES * STAGES; +const C_OFFSET: usize = B_OFFSET + B_BYTES * STAGES; +const C_STAGE_BYTES: usize = 128 * 32 * 2; +const BARRIER_OFFSET: usize = C_OFFSET + 2 * C_STAGE_BYTES; +pub const DYNAMIC_SMEM_BYTES: usize = BARRIER_OFFSET + 128; +const _: () = assert!(DYNAMIC_SMEM_BYTES == 211_072); + +type MainloopPipeline = Sm100TmaMmaPipeline<5, 77_824>; +type AccumulatorPipeline = Sm100AccumulatorPipeline<4>; +type TiledMma = Sm100TiledMma; + +// Only the allocation rendezvous uses an ordinary cluster barrier. Operand +// and accumulator lifetimes are described by the CuTe pipelines below. +#[inline(always)] +unsafe fn wait_deallocation(bar: *const Barrier) { + unsafe { while !mbarrier_try_wait_parity_cluster(bar, 0) {} } +} + +#[cuda_module] +pub mod kernels { + use super::*; + + /// Requires a (2,1,1) cluster launch and the four tensor maps documented + /// in README.md. All pointers must remain live through stream completion. + #[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] + pub fn fp16_gemm_256x352( + a: *const TmaDesc, + b0: *const TmaDesc, + b1: *const TmaDesc, + c: *const TmaDesc, + bias: *const f16, + m: u64, + n: u64, + k: u64, + has_bias: u64, + ) { + // The pinned backend's direct-launch scalar ABI uses 64-bit carriers. + // Host validation bounds all dimensions before these device casts. + 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::block_rank(); + let smem = DynamicSharedArray::::get_raw(); + // ABI: A[5], B[5], C[2], full[5], empty[5], acc_empty, + // acc_full, dealloc, tmem_token. All operand bases are B128 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 = Sm100TmaStorePipeline::<2>::new(); + if warp_id == 4 { + prefetch_tma_descriptor(a.cast()); + prefetch_tma_descriptor(b0.cast()); + prefetch_tma_descriptor(b1.cast()); + prefetch_tma_descriptor(c.cast()); + } + if warp_id == 0 && warp::is_elected_sync(u32::MAX) { + mbarrier_init(dealloc, 32); + accumulator.init(); + pipeline.init(); + } + fence_mbarrier_init_release_cluster(); + cluster::barrier_cluster_arrive_relaxed(); + + // Default source scheduler: M-major, one logical tile per cluster. + let mtiles = m / 256; + let total_tiles = mtiles * ((n + 351) / 352); + let mut tile = thread::blockIdx_x() / 2; + let tile_stride = thread::gridDim_x() / 2; + let ktiles = (k + 63) / 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(stage as u32); + } + // Native CuTe cluster TMA performs its own warp election. + // Every producer lane must participate in each copy. + let barrier = pipeline.full_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_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) { + // The paired tile owns the A collector lifetime for + // each of the four K16 instructions in this stage. + 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 { + Sm100Tmem::<512>::allocate(tmem_token); + } + 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_proxy_async_shared_cta(); + 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_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_cta_sync(3, 32); + let peer = cluster::map_shared_rank(dealloc, rank ^ 1); + mbarrier_arrive_cluster(peer as u64); + wait_deallocation(dealloc); + Sm100Tmem::<512>::deallocate(tmem); + } + } + } + } +} diff --git a/cuteir/examples/fp16_gemm_256x352_cute/src/main.rs b/cuteir/examples/fp16_gemm_256x352_cute/src/main.rs new file mode 100644 index 0000000000..37bba38296 --- /dev/null +++ b/cuteir/examples/fp16_gemm_256x352_cute/src/main.rs @@ -0,0 +1,559 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * SPDX-License-Identifier: LicenseRef-NvidiaProprietary + */ +#![feature(f16)] + +mod kernel; + +use cuda_core::{CudaContext, CudaStream, DeviceBuffer, LaunchConfig1D, sys}; +use cute_rs::tma::{TmaDesc, make_tma_desc_2d}; +use serde_json::json; +use std::{ + error::Error, + path::{Path, PathBuf}, + sync::Arc, +}; + +type Result = std::result::Result>; + +#[derive(Clone, Copy, PartialEq, Eq)] +enum TimingMode { + Direct, + Graph, +} + +impl TimingMode { + fn name(self) -> &'static str { + match self { + Self::Direct => "direct", + Self::Graph => "graph", + } + } + + fn tool(self) -> &'static str { + match self { + Self::Direct => "cuda-events-single-kernel", + Self::Graph => "cuda-graph-events-per-kernel", + } + } +} + +struct Args { + m: u32, + n: u32, + k: u32, + has_bias: bool, + warmup: usize, + iters: usize, + graph_launches: usize, + timing_mode: TimingMode, + clusters: Option, + input_dir: Option, + output: Option, + json: Option, + skip_cpu_check: bool, +} + +impl Args { + fn parse() -> Result { + let mut args = Self { + 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 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 [--timing-mode direct|graph] [--graph-launches 10] [--clusters COUNT] [--input-dir DIR] [--output C.f16] [--json RESULT.json]\n [--skip-cpu-check] (requires --input-dir and --output; external oracle mode)\nTiming defaults to graph; direct times exactly one kernel launch per event pair.\nDirect --warmup 0 --iters 1 launches once total and verifies that timed output.\n--graph-launches applies only to graph timing.\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, + _ => { + 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()?, + "--timing-mode" => { + 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()), + _ => return Err(format!("unknown argument {key}").into()), + } + } + } + } + if args.m == 0 + || args.n == 0 + || args.k == 0 + || args.m % 256 != 0 + || args.n % 8 != 0 + || args.k % 8 != 0 + { + 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. Bounds above also make n+351 and + // k+63 safe in the kernel's u32 ceil divisions. + 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) + } +} + +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. +struct KernelGraph { + executable: sys::CUgraphExec, + stream: Arc, +} + +impl KernelGraph { + 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), + }) + } + + 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); + } + } + } +} + +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 + .chunks_exact(2) + .map(|v| u16::from_le_bytes([v[0], v[1]])) + .collect()) +} + +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(()) +} + +// Reproducible integer inputs use the oracle's exact FP16 domain [-2,2). +// compare.py supplies the oracle's actual seeded PyTorch buffers instead. +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() +} + +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(()) +} + +fn main() -> Result<()> { + let args = Args::parse()?; + 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])?; + // The same layout types describe the host TMA box/swizzle and the + // device shared-memory tensors, so these contracts cannot drift apart. + let a_desc = make_tma_desc_2d::( + 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 = make_tma_desc_2d::( + 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 = make_tma_desc_2d::( + 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 = make_tma_desc_2d::( + c_dev.cu_deviceptr() as *mut core::ffi::c_void, + u64::from(args.m), + u64::from(args.n), + u64::from(args.n), + )?; + // Each allocation contains exactly one descriptor's 128 encoded bytes. + // CUDA device allocations satisfy TmaDesc's 64-byte alignment, and the + // buffers remain alive until every direct launch or graph replay completes. + let a_desc_dev = DeviceBuffer::from_host(&stream, &a_desc.bytes)?; + let b0_desc_dev = DeviceBuffer::from_host(&stream, &b0_desc.bytes)?; + let b1_desc_dev = DeviceBuffer::from_host(&stream, &b1_desc.bytes)?; + let c_desc_dev = DeviceBuffer::from_host(&stream, &c_desc.bytes)?; + // 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 as u32, + kernel::DYNAMIC_SMEM_BYTES as u32, + ))?; + let capacity = probe.function().max_active_clusters( + (blocks, 1, 1), + (kernel::THREADS as u32, 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 as u32, + kernel::DYNAMIC_SMEM_BYTES as u32, + ))?; + let run = || -> Result<()> { + // Descriptors address live allocations; the prepared contract + // fixes block, cluster and shared-memory shape. Persistent tiles are + // disjoint, and the tensors satisfy the kernel's shape constraints. + module.fp16_gemm_256x352( + &stream, + &launch, + a_desc_dev.cu_deviceptr() as *const TmaDesc, + b0_desc_dev.cu_deviceptr() as *const TmaDesc, + b1_desc_dev.cu_deviceptr() as *const TmaDesc, + c_desc_dev.cu_deviceptr() as *const TmaDesc, + 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": "cuda-oxide", "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/cuteir/examples/fp16_gemm_256x352_cute/validate.py b/cuteir/examples/fp16_gemm_256x352_cute/validate.py new file mode 100644 index 0000000000..a3e14a918e --- /dev/null +++ b/cuteir/examples/fp16_gemm_256x352_cute/validate.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# 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), "--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), "--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()) From f766eccd6194245c7ebe1f7cdfae4d2498cd5dfe Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Wed, 16 Sep 2026 14:16:20 +0000 Subject: [PATCH 07/99] chore(vscode): include SM100 CuTe GEMM in rust-analyzer projects --- .vscode/settings.json | 1 + 1 file changed, 1 insertion(+) diff --git a/.vscode/settings.json b/.vscode/settings.json index 12d1042dc6..719501d548 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -24,6 +24,7 @@ "./cuteir/examples/elementwise_cute/Cargo.toml", "./cuteir/examples/blockscale_gemm_cute/Cargo.toml", "./cuteir/examples/nvfp4_gemv_cute/Cargo.toml", + "./cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml", "./crates/rustc-codegen-cuda/examples/debug_pointer_locals/Cargo.toml", "./crates/rustc-codegen-cuda/examples/gemm_sol_final/Cargo.toml", ], From e1644effebe748449867d9f4b3ecab28ca195927 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Wed, 16 Sep 2026 20:00:56 +0530 Subject: [PATCH 08/99] fix(cuteir): correct SM100 GEMM license metadata Use Apache-2.0 headers for the Rust kernel and host runner and align the example package license. This change only updates license comments and metadata. Signed-off-by: nihalpasham --- cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml | 2 +- .../examples/fp16_gemm_256x352_cute/src/kernel.rs | 13 ++++--------- cuteir/examples/fp16_gemm_256x352_cute/src/main.rs | 5 +++-- 3 files changed, 8 insertions(+), 12 deletions(-) diff --git a/cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml b/cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml index 18a2bddb4a..34adaaef83 100644 --- a/cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml +++ b/cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml @@ -2,7 +2,7 @@ name = "fp16_gemm_256x352_cute" version = "0.1.0" edition = "2024" -license = "LicenseRef-NvidiaProprietary" +license = "Apache-2.0" description = "SM100 two-CTA persistent FP16 GEMM with a 256x352 output tile" [workspace] diff --git a/cuteir/examples/fp16_gemm_256x352_cute/src/kernel.rs b/cuteir/examples/fp16_gemm_256x352_cute/src/kernel.rs index 1e0b3d7a86..ae013b29c4 100644 --- a/cuteir/examples/fp16_gemm_256x352_cute/src/kernel.rs +++ b/cuteir/examples/fp16_gemm_256x352_cute/src/kernel.rs @@ -1,12 +1,7 @@ -// {$nv-internal-release file} -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: LicenseRef-NvidiaProprietary -// -// Rust port of NVIDIA DKG fp16_gemm_3_256x352.py (revision b3f067d4). -// NVIDIA CORPORATION, its -// affiliates and licensors retain all intellectual property and proprietary -// rights in and to this material and modifications thereto. Use, reproduction, -// disclosure or distribution requires an express NVIDIA license agreement. +/* + * 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 diff --git a/cuteir/examples/fp16_gemm_256x352_cute/src/main.rs b/cuteir/examples/fp16_gemm_256x352_cute/src/main.rs index 37bba38296..c456f7baf6 100644 --- a/cuteir/examples/fp16_gemm_256x352_cute/src/main.rs +++ b/cuteir/examples/fp16_gemm_256x352_cute/src/main.rs @@ -1,7 +1,8 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. - * SPDX-License-Identifier: LicenseRef-NvidiaProprietary + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 */ + #![feature(f16)] mod kernel; From 4368493caab29454e59bca7432e80796928eba1b Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Fri, 18 Sep 2026 13:08:54 +0000 Subject: [PATCH 09/99] feat(cuteir): pass TMA descriptors as grid constants Carry grid-constant parameter size and alignment through CUTLASS MLIR export and cubin ABI validation. Pass the SM100 GEMM descriptors directly with the launch and preserve assertion behavior after rebasing onto main. Run the local CuTe DSL comparison by default, with standalone execution available through --rust-only. Keep the local Python reference assets untracked and document the setup and timing differences. Update CuTe examples for the shared host crates. Simplify documentation, fix test compilation errors and Clippy warnings, and include the new crates in the book inventory and CI test matrix. Validation: 5,636 workspace tests, 34 doctests, strict root/backend Clippy, all-feature analyzer checks, formatting, rustdoc and book warning checks, and repository consistency guards passed. Signed-off-by: nihalpasham --- .github/workflows/unit-tests.yml | 7 + Cargo.lock | 8 +- Justfile | 3 +- crates/cargo-oxide/src/commands/tests.rs | 2 + .../cuda-oxide-codegen/src/cutlass_backend.rs | 255 +++++++++++++++--- .../src/cutlass_compiler.rs | 11 +- crates/cuda-oxide-codegen/src/pipeline.rs | 12 +- crates/cuda-oxide-codegen/src/prep.rs | 21 +- .../src/tests/cutlass/grid_constant.rs | 175 ++++++++++++ .../src/tests/cutlass/tcgen05.rs | 13 +- crates/cuda-oxide-mlir-export/Cargo.toml | 1 + crates/cuda-oxide-mlir-export/README.md | 60 +++-- crates/cuda-oxide-mlir-export/src/cute.rs | 12 +- .../cuda-oxide-mlir-export/src/cute_gemm.rs | 6 +- .../cuda-oxide-mlir-export/src/cute_gemv.rs | 9 +- .../cuda-oxide-mlir-export/src/cute_sm100.rs | 16 +- .../src/grid_constant.rs | 151 +++++++++++ crates/cuda-oxide-mlir-export/src/lib.rs | 6 +- crates/cuda-oxide-mlir-export/src/mir_core.rs | 47 +++- .../cuda-oxide-mlir-export/src/mir_memory.rs | 64 ++++- .../src/nvvm_tcgen05.rs | 16 +- crates/cuda-oxide-mlir-export/src/profile.rs | 10 +- .../tests/grid_constant.rs | 190 +++++++++++++ .../mir-importer/src/translator/cute/mod.rs | 46 ++-- .../mir-lower/src/convert/ops/control_flow.rs | 3 +- crates/rustc-codegen-cuda/Cargo.lock | 1 + .../compiler/architecture-overview.md | 7 +- cuteir/README.md | 31 +-- cuteir/cute-rs/src/sm100.rs | 67 +++-- cuteir/cute-rs/src/tma.rs | 38 ++- cuteir/dialect-cute/src/lib.rs | 19 +- cuteir/dialect-cute/src/sm100_ops.rs | 8 +- cuteir/dialect-cute/src/types.rs | 24 +- cuteir/dialect-cute/src/verify.rs | 8 +- cuteir/docs/translation.md | 42 ++- .../examples/blockscale_gemm_cute/Cargo.lock | 193 ++++++++++++- .../examples/blockscale_gemm_cute/Cargo.toml | 2 +- cuteir/examples/elementwise_cute/Cargo.lock | 193 ++++++++++++- cuteir/examples/elementwise_cute/Cargo.toml | 2 +- cuteir/examples/elementwise_cute/src/main.rs | 3 +- .../fp16_gemm_256x352_cute/.gitignore | 1 + .../fp16_gemm_256x352_cute/Cargo.lock | 193 ++++++++++++- .../fp16_gemm_256x352_cute/Cargo.toml | 2 +- .../examples/fp16_gemm_256x352_cute/README.md | 111 ++++++-- .../fp16_gemm_256x352_cute/src/kernel.rs | 45 ++-- .../fp16_gemm_256x352_cute/src/main.rs | 162 ++++++++--- .../fp16_gemm_256x352_cute/validate.py | 6 +- cuteir/examples/nvfp4_gemv_cute/Cargo.lock | 193 ++++++++++++- cuteir/examples/nvfp4_gemv_cute/Cargo.toml | 2 +- cuteir/examples/nvfp4_gemv_cute/src/main.rs | 3 +- scripts/check-shared-crate-pin.sh | 8 +- 51 files changed, 2110 insertions(+), 398 deletions(-) create mode 100644 crates/cuda-oxide-codegen/src/tests/cutlass/grid_constant.rs create mode 100644 crates/cuda-oxide-mlir-export/src/grid_constant.rs create mode 100644 crates/cuda-oxide-mlir-export/tests/grid_constant.rs diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml index 35a76d73d0..d14b0fa6eb 100644 --- a/.github/workflows/unit-tests.yml +++ b/.github/workflows/unit-tests.yml @@ -58,6 +58,13 @@ jobs: - package: dialect-mir - package: dialect-nvvm - package: mir-importer + # CuTe types, semantic IR, and MLIR export need no CUDA toolkit. + # cute-rs uses its default device API; its host helpers are optional. + - package: cute-layout + - package: cute-rs + - package: dialect-cute + - package: pliron-mlir-export + - package: cuda-oxide-mlir-export # Includes a real llc -> ptxas spine-kernel validation. - package: cuda-oxide-codegen needs_cuda: true diff --git a/Cargo.lock b/Cargo.lock index 0379b70843..dbe265ec77 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -518,9 +518,10 @@ dependencies = [ "pliron", "pliron-derive", "pliron-mlir-export", + "reserved-oxide-symbols", "serde", "serde_json", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -541,6 +542,7 @@ dependencies = [ "trybuild", ] +[[package]] name = "dashmap" version = "6.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -565,7 +567,7 @@ dependencies = [ "once_cell", "pliron", "pliron-derive", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -1236,7 +1238,7 @@ version = "0.1.0" dependencies = [ "pliron", "pliron-derive", - "thiserror", + "thiserror 2.0.18", ] [[package]] diff --git a/Justfile b/Justfile index b00f1a6651..7dc9e9f3ae 100644 --- a/Justfile +++ b/Justfile @@ -64,7 +64,8 @@ test: -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 cute-rs -p dialect-cute \ + -p pliron-mlir-export -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 diff --git a/crates/cargo-oxide/src/commands/tests.rs b/crates/cargo-oxide/src/commands/tests.rs index 48986e6e9d..9594efc0d7 100644 --- a/crates/cargo-oxide/src/commands/tests.rs +++ b/crates/cargo-oxide/src/commands/tests.rs @@ -2991,6 +2991,7 @@ fn passthrough_fingerprint_tracks_the_managed_cutlass_library_digest() { 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()), @@ -3042,6 +3043,7 @@ fn passthrough_fingerprint_tracks_explicit_cutlass_library_content() { 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()), diff --git a/crates/cuda-oxide-codegen/src/cutlass_backend.rs b/crates/cuda-oxide-codegen/src/cutlass_backend.rs index 57b399abd6..50ea536096 100644 --- a/crates/cuda-oxide-codegen/src/cutlass_backend.rs +++ b/crates/cuda-oxide-codegen/src/cutlass_backend.rs @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! Executable Backend-B boundary for the official CUTLASS compiler library. +//! Compile MLIR with the pinned CUTLASS library and validate its CUDA image. use crate::cutlass_compiler::CutlassCompilerLibrary; use crate::error::PipelineError; @@ -47,7 +47,11 @@ mod sm100_tests; #[path = "tests/cutlass/tcgen05.rs"] mod tcgen05_tests; -/// External input needed to execute the official CUTLASS Backend B. +#[cfg(test)] +#[path = "tests/cutlass/grid_constant.rs"] +mod grid_constant_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 @@ -70,7 +74,7 @@ impl CutlassBackendConfig { /// Device compiler selected after the shared MIR/CuTe preparation stages. #[derive(Clone, Debug, Default, PartialEq, Eq)] pub enum DeviceBackend { - /// Ordinary CUDA Oxide MIR/NVVM/LLVM backend inherited from main. + /// CUDA Oxide MIR/NVVM/LLVM backend. #[default] Native, /// Official CUTLASS compiler-library path, producing a CUDA image. @@ -86,17 +90,28 @@ pub(crate) struct CutlassRunOutput { enum KernelParameterKind { GlobalPointer64, Scalar64, + GridConstant { size: u16, alignment: u16 }, } impl KernelParameterKind { const fn size(self) -> u16 { - 8 + match self { + Self::GlobalPointer64 | Self::Scalar64 => 8, + Self::GridConstant { size, .. } => size, + } + } + + const fn alignment(self) -> u16 { + match self { + Self::GlobalPointer64 | Self::Scalar64 => 8, + Self::GridConstant { alignment, .. } => alignment, + } } const fn cuda_parameter_space(self) -> u32 { match self { Self::GlobalPointer64 => 5, - Self::Scalar64 => 0, + Self::Scalar64 | Self::GridConstant { .. } => 0, } } } @@ -203,17 +218,16 @@ fn add_execution_manifest( Ok(output) } -/// Recover the exact direct-launch ABI that the pinned export profile writes -/// on each kernel. Backend B currently flattens every Rust slice into a -/// `(pointer, i64)` pair and every remaining scalar into `i64`; accepting any -/// other type here would require a corresponding host-launch ABI decision. +/// 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 `i64` scalars. 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\"() <{function_type = ("; - const RETURN_AND_NAME: &str = ") -> (), sym_name = \""; - + const PREFIX: &str = "\"func.func\"() <{"; let mut abis = BTreeMap::new(); for kernel in expected_kernels { if kernel.contains(['"', '\\']) { @@ -222,18 +236,13 @@ fn expected_kernel_abis( if abis.contains_key(kernel) { return Err(format!("duplicate expected kernel {kernel:?}")); } - let name_marker = format!("{RETURN_AND_NAME}{kernel}\""); - // The generic printer may append whitespace after the region opener; - // match the stable function/name portion instead of depending on it. - let mut signatures = rendered_mlir + let name_marker = format!("sym_name = \"{kernel}\""); + let signatures = rendered_mlir .lines() .filter_map(|line| { - let signature = line.trim().strip_prefix(PREFIX)?; - let end = signature.find(&name_marker)?; - signature[end + name_marker.len()..] - .trim_start() - .starts_with("}> ({") - .then_some(&signature[..end]) + let properties = line.trim().strip_prefix(PREFIX)?; + let (properties, _) = properties.split_once("}> ({")?; + properties.contains(&name_marker).then_some(properties) }) .collect::>(); if signatures.len() != 1 { @@ -242,26 +251,172 @@ fn expected_kernel_abis( signatures.len() )); } - let signature = signatures.pop().expect("length checked"); - let parameters = if signature.trim().is_empty() { - Vec::new() - } else { - signature - .split(',') - .map(|parameter| match parameter.trim() { - "!llvm.ptr" => Ok(KernelParameterKind::GlobalPointer64), - "i64" => Ok(KernelParameterKind::Scalar64), - other => Err(format!( - "kernel {kernel:?} has unsupported flattened parameter type {other:?}; only !llvm.ptr and i64 have a pinned direct-launch ABI" - )), - }) - .collect::, _>>()? - }; + 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), + "i64" => Ok(KernelParameterKind::Scalar64), + other => Err(format!( + "unsupported flattened parameter type {other:?}; only !llvm.ptr and i64 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 @@ -552,10 +707,11 @@ fn validate_cuda_cubin( Ok(()) } -/// Validate the post-ptxas direct-launch parameter layout recorded in CUDA's -/// per-kernel `.nv.info` section. These record tags and bit fields are part of -/// the fingerprint-pinned CUTLASS 4.7 / CUDA 13.x output contract; unknown or -/// malformed records are rejected instead of being guessed through. +/// 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], @@ -666,6 +822,11 @@ fn validate_kernel_parameter_abi( 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}"))?; @@ -876,19 +1037,27 @@ mod tests { 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(&((ordinal * 8) 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(&((parameters.len() * 8) as u16).to_le_bytes()); + info.extend_from_slice(&offset.to_le_bytes()); info } diff --git a/crates/cuda-oxide-codegen/src/cutlass_compiler.rs b/crates/cuda-oxide-codegen/src/cutlass_compiler.rs index 13e2d0f5ff..fdc0ee02f6 100644 --- a/crates/cuda-oxide-codegen/src/cutlass_compiler.rs +++ b/crates/cuda-oxide-codegen/src/cutlass_compiler.rs @@ -5,12 +5,11 @@ //! Runtime-loaded bindings for the CUTLASS 4.7 compiler C API. //! -//! This module intentionally exposes one narrow operation: compile textual -//! `PreCompiledMlir` through a pinned CuTe pipeline to the C API's -//! `ObjectArtifact`. Despite containing device code, that artifact is a host -//! ELF object, not a standalone CUDA image. The -//! official API owns all compiler and artifact handles, so the wrappers below -//! release them through the matching C destructors. +//! 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}; diff --git a/crates/cuda-oxide-codegen/src/pipeline.rs b/crates/cuda-oxide-codegen/src/pipeline.rs index be49b66b45..a576f8cffa 100644 --- a/crates/cuda-oxide-codegen/src/pipeline.rs +++ b/crates/cuda-oxide-codegen/src/pipeline.rs @@ -187,7 +187,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, @@ -280,9 +280,8 @@ pub fn compile_translated_module( .emit(format!("{}", module.deref(ctx).disp(ctx))); } - // This is the shared CuTe backend seam. Validate the complete semantic - // story once, before either a native lowering or an external compiler - // translation is allowed to consume it. + // Check tensor origins, pipeline state, and operation ordering before + // either backend lowers the CuTe operations. dialect_cute::verify::verify_cute_semantics(ctx, module).map_err(|error| { PipelineError::Lowering(format!( "dialect-cute semantic verification failed: {error}" @@ -294,9 +293,8 @@ pub fn compile_translated_module( // materialization fails with its explicit-target requirement). materialize_iket(ctx, module, backend.target_arch.as_deref(), &backend.iket)?; - // Observe or compile the shared prepared semantic module before the - // ordinary inherited backend begins MIR/NVVM/LLVM lowering. Builds without - // an observation or CUTLASS selection do no translation-backend work. + // 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, backend)?) } else { diff --git a/crates/cuda-oxide-codegen/src/prep.rs b/crates/cuda-oxide-codegen/src/prep.rs index 81f642f2d5..74a5c6388a 100644 --- a/crates/cuda-oxide-codegen/src/prep.rs +++ b/crates/cuda-oxide-codegen/src/prep.rs @@ -196,14 +196,11 @@ pub fn prepare_mir_module( )?; if has_high_level_cute_operations { - // A reborrow of a local CuTe carrier binding (`&mut tensor`) imports - // as a pointer-kind-only `mir.cast` on the alloca slot pointer. That - // cast is a non-promotable use for mem2reg, but ghost CuTe values are - // forbidden from living in memory at all, so the kind refinement - // carries no ABI meaning for them. Once the first mem2reg pass has - // dissolved the pointer slots, the cast's remaining uses are plain - // memory addresses; fold it back onto the slot pointer so the second - // pass can dissolve the ghost slot itself. + // 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")?; @@ -335,10 +332,10 @@ fn fold_ghost_slot_kind_only_casts(ctx: &mut Context, root: Ptr) { 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) { - if is_ghost_slot_kind_only_cast(ctx, &cast) { - foldable.push(child); - } + if let Some(cast) = Operation::get_op::(child, ctx) + && is_ghost_slot_kind_only_cast(ctx, &cast) + { + foldable.push(child); } } } diff --git a/crates/cuda-oxide-codegen/src/tests/cutlass/grid_constant.rs b/crates/cuda-oxide-codegen/src/tests/cutlass/grid_constant.rs new file mode 100644 index 0000000000..159eb23a9d --- /dev/null +++ b/crates/cuda-oxide-codegen/src/tests/cutlass/grid_constant.rs @@ -0,0 +1,175 @@ +/* + * 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, + } +} + +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/crates/cuda-oxide-codegen/src/tests/cutlass/tcgen05.rs b/crates/cuda-oxide-codegen/src/tests/cutlass/tcgen05.rs index 5936d30e9a..78e7bbdad2 100644 --- a/crates/cuda-oxide-codegen/src/tests/cutlass/tcgen05.rs +++ b/crates/cuda-oxide-codegen/src/tests/cutlass/tcgen05.rs @@ -120,10 +120,10 @@ fn cutlass_tcgen05_pack_compiles() { Some("v1:i0763"), ); let mma = Tcgen05MmaOp::new(op); - mma.set_attr_nvvm_tcgen05_mma_form(&mut ctx, Tcgen05MmaFormAttr::Shared); - mma.set_attr_nvvm_tcgen05_mma_kind(&mut ctx, Tcgen05MmaKindAttr::F16); - mma.set_attr_nvvm_tcgen05_mma_cta_group(&mut ctx, Tcgen05MmaCtaGroupAttr::Cg2); - mma.set_attr_nvvm_tcgen05_mma_collector_a(&mut ctx, collector); + 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, @@ -236,7 +236,7 @@ fn cutlass_cluster_tma_pack_compiles() { (ClusterBarrierModeAttr::Wait, "v1:i0281"), ] { let op = append::(&mut ctx, block, vec![], vec![], Some(marker)); - ClusterBarrierOp::new(op).set_attr_nvvm_cluster_barrier_mode(&mut ctx, mode); + ClusterBarrierOp::new(op).set_attr_nvvm_cluster_barrier_mode(&ctx, mode); } append::( &mut ctx, @@ -310,8 +310,7 @@ fn cutlass_cluster_tma_pack_compiles() { vec![u64.into()], None, ); - MirCastOp::new(remote_address) - .set_attr_cast_kind(&mut ctx, MirCastKindAttr::PointerExposeAddress); + 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![], vec![], None); diff --git a/crates/cuda-oxide-mlir-export/Cargo.toml b/crates/cuda-oxide-mlir-export/Cargo.toml index e002335f5c..5ca4d1c03b 100644 --- a/crates/cuda-oxide-mlir-export/Cargo.toml +++ b/crates/cuda-oxide-mlir-export/Cargo.toml @@ -15,6 +15,7 @@ dialect-mir = { workspace = true } dialect-nvvm = { workspace = true } dialect-cute = { workspace = true } llvm-export = { workspace = true } +reserved-oxide-symbols = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } thiserror = { workspace = true } diff --git a/crates/cuda-oxide-mlir-export/README.md b/crates/cuda-oxide-mlir-export/README.md index 5299d09e55..c5f8bc6f41 100644 --- a/crates/cuda-oxide-mlir-export/README.md +++ b/crates/cuda-oxide-mlir-export/README.md @@ -1,9 +1,7 @@ # CUDA Oxide MLIR export -This crate contains CUDA Oxide's target-specific mappings for -`pliron-mlir-export`. It is the CUTLASS translation backend. The generic -crate builds typed, deterministic MLIR text; this crate explains what each -CUDA Oxide operation means to one pinned MLIR consumer. +This crate maps CUDA Oxide operations to MLIR accepted by the pinned CUTLASS +compiler. It uses `pliron-mlir-export` to build and render the MLIR text: ```text pliron-mlir-export this crate @@ -12,10 +10,10 @@ syntax tree + renderer + builtin / MIR / NVVM / CuTe mappings registry + diagnostics pinned consumer profiles ``` -The split is deliberate. The generic crate can translate any Pliron dialect. -This crate decides what CUDA Oxide operations mean in one exact MLIR consumer. +The generic crate can translate any Pliron dialect. CUDA-specific mappings +and compiler-version requirements live here. -The first profile targets the official CUTLASS 4.7 compiler library. CUDA +The profile targets the official CUTLASS 4.7 compiler library. CUDA Oxide installs the exact NVIDIA release archive, verifies both the archive and library digests, and loads that versioned library explicitly: @@ -35,12 +33,11 @@ shared high-level Pliron module official CUTLASS 4.7 compiler ``` -The elementwise, NVFP4 GEMV, and block-scaled GEMM mapping packs translate the -live post-preparation module directly. The pinned library compiles that MLIR to -a validated cubin, which is embedded and launched through CUDA Oxide's ordinary -artifact loader. Set `CUDA_OXIDE_MLIR_OUTPUT=` to retain the exact textual -module for inspection; select `CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir` to compile -it. +The mapping packs translate CuTe operations after shared MIR preparation. +The pinned library compiles the resulting MLIR, and the backend extracts and +validates its cubin. CUDA Oxide embeds and launches that cubin through its +ordinary artifact loader. Set `CUDA_OXIDE_MLIR_OUTPUT=` to keep the MLIR +text for inspection; 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, @@ -50,13 +47,14 @@ 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. -## The first mapping pack +## Scalar code and control flow -The first pack covers the ordinary scalar code around a CuTe kernel: +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 @@ -121,6 +119,24 @@ 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 @@ -128,10 +144,9 @@ gpu_kernel = "true" → cute.kernel anything else → translation error ``` -## The first CuTe slice +## Tensor operations -The elementwise path keeps the useful layout explanation instead of reducing -everything to pointer arithmetic first: +The elementwise path keeps tensor layouts visible until CUTLASS lowers them: ```text make tensor @@ -147,11 +162,10 @@ pick this thread's tile └── edge tile ──► checked scalar load/store ``` -The profile recognizes the corresponding `cute.tensor_*` operations and maps -them to CUTLASS's `cute` and `cute_nvgpu` operations. Before translation, it -runs the shared backend-neutral whole-module verifier directly. The exporter -does not invoke native CuTe expansion or use native generated-intrinsic markers -as an intermediate representation. +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: diff --git a/crates/cuda-oxide-mlir-export/src/cute.rs b/crates/cuda-oxide-mlir-export/src/cute.rs index e20b82c30d..1871774afe 100644 --- a/crates/cuda-oxide-mlir-export/src/cute.rs +++ b/crates/cuda-oxide-mlir-export/src/cute.rs @@ -3,10 +3,9 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! The first semantic CuTe mapping slice for the official CUTLASS `cute` -//! dialect. +//! Map tensor views and copies to the official CUTLASS `cute` dialect. //! -//! The split is deliberate: +//! Full tiles use CuTe operations; edge tiles use checked scalar accesses: //! //! ```text //! tensor_make -> zipped_divide -> slice -> full-tile copy @@ -16,10 +15,9 @@ //! ordinary arith + LLVM, using verified producer provenance //! ``` //! -//! CUTLASS has first-class operations for the view transformations and bulk copy. -//! It does not have operations carrying the original logical length or the -//! absolute-index tail-store contract, so inventing similarly named target -//! operations would lose semantics rather than preserve them. +//! 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 dialect_cute::{ attributes::{ diff --git a/crates/cuda-oxide-mlir-export/src/cute_gemm.rs b/crates/cuda-oxide-mlir-export/src/cute_gemm.rs index f7a090f5e4..735b80812d 100644 --- a/crates/cuda-oxide-mlir-export/src/cute_gemm.rs +++ b/crates/cuda-oxide-mlir-export/src/cute_gemm.rs @@ -6,9 +6,9 @@ //! SM120 block-scaled GEMM mappings for the official CUTLASS 4.7 compiler. //! //! The source dialect keeps scheduler, pipeline, TMA, MMA, and epilogue -//! contracts visible. This pack erases only compiler-only handles and lowers -//! their runtime leaves to operations accepted by the public PRE_COMPILED -//! pipeline in `libCutlassCompiler.so`. +//! operations visible. This pack removes compile-time handles and maps their +//! runtime work to operations accepted by the public PRE_COMPILED pipeline +//! in `libCutlassCompiler.so`. use dialect_cute::{ attributes::{ diff --git a/crates/cuda-oxide-mlir-export/src/cute_gemv.rs b/crates/cuda-oxide-mlir-export/src/cute_gemv.rs index 5eadd84bda..866c1d0822 100644 --- a/crates/cuda-oxide-mlir-export/src/cute_gemv.rs +++ b/crates/cuda-oxide-mlir-export/src/cute_gemv.rs @@ -5,9 +5,8 @@ //! NVFP4 GEMV mappings for the official CUTLASS full-CuTe MLIR profile. //! -//! The source views are compiler-only values, but the target carriers are -//! ordinary LLVM aggregates. Keeping the runtime pointer and shape fields in -//! SSA gives each selection operation a real target dataflow edge: +//! 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> @@ -19,8 +18,8 @@ //! load_vec -> cvt_fpext -> ordered f32 dot //! ``` //! -//! This translation consumes the shared semantic operations directly and -//! manufactures only target operations owned by CUTLASS. +//! The mapping translates CuTe operations directly into operations accepted +//! by CUTLASS. use crate::cute::{TensorViewTypeTranslation, register_cute_elementwise_mappings}; diff --git a/crates/cuda-oxide-mlir-export/src/cute_sm100.rs b/crates/cuda-oxide-mlir-export/src/cute_sm100.rs index a4fe9c30bd..a278d6ba4b 100644 --- a/crates/cuda-oxide-mlir-export/src/cute_sm100.rs +++ b/crates/cuda-oxide-mlir-export/src/cute_sm100.rs @@ -3,14 +3,18 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! SM100 CuTe semantic plans lowered by the pinned CUTLASS 4.7 compiler. +//! 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. The public 4.7 `sm100.mma` atom exposes no A -//! collector field (its fields are accumulate, negate A/B and disable lanes). -//! The collector-aware MMA instruction therefore uses an NVVM architecture leaf: -//! its descriptors come from CuTe layouts and its shape/collector sequence come -//! from the verified tiled-MMA plan, never raw Rust instruction 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 +//! ``` use dialect_cute::{ attributes::{CuteComposedLayoutAttr, CuteTmaStorePipelineAttr}, diff --git a/crates/cuda-oxide-mlir-export/src/grid_constant.rs b/crates/cuda-oxide-mlir-export/src/grid_constant.rs new file mode 100644 index 0000000000..87af4b7102 --- /dev/null +++ b/crates/cuda-oxide-mlir-export/src/grid_constant.rs @@ -0,0 +1,151 @@ +/* + * 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 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 pliron_mlir_export::{MlirAttribute, MlirOperation, MlirType}; +use reserved_oxide_symbols::{ + MIR_GRID_CONSTANT_ALIGN_ATTR_PREFIX, MIR_GRID_CONSTANT_POINTEE_ATTR_PREFIX, +}; + +pub(crate) struct Parameter { + size: u64, + alignment: u64, +} + +impl Parameter { + 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/crates/cuda-oxide-mlir-export/src/lib.rs b/crates/cuda-oxide-mlir-export/src/lib.rs index 3f05296b89..27ca96bc84 100644 --- a/crates/cuda-oxide-mlir-export/src/lib.rs +++ b/crates/cuda-oxide-mlir-export/src/lib.rs @@ -5,14 +5,14 @@ //! CUDA Oxide mappings for textual MLIR export. //! -//! The generic exporter answers "how do we build typed, deterministic MLIR -//! text?" This crate answers "what does this CUDA Oxide operation become for -//! this exact consumer?" +//! The generic exporter builds typed MLIR text. This crate maps CUDA Oxide +//! operations to the types and operations accepted by the pinned CUTLASS compiler. mod cute; mod cute_gemm; mod cute_gemv; mod cute_sm100; +mod grid_constant; mod mir_core; mod mir_memory; mod nvvm_cluster; diff --git a/crates/cuda-oxide-mlir-export/src/mir_core.rs b/crates/cuda-oxide-mlir-export/src/mir_core.rs index 1516f98d23..82c703b612 100644 --- a/crates/cuda-oxide-mlir-export/src/mir_core.rs +++ b/crates/cuda-oxide-mlir-export/src/mir_core.rs @@ -10,7 +10,9 @@ //! signed or unsigned `arith` operation before that information disappears. use dialect_mir::{ - attributes::{MirCastKindAttr, MirFP16Attr, MirPointerKindAuthorityAttr}, + attributes::{ + MirCastKindAttr, MirFP16Attr, MirPointerKindAuthorityAttr, ReferenceParamValidityAttr, + }, ops::{ MirAddOp, MirBitAndOp, MirBitOrOp, MirBitXorOp, MirCastOp, MirCondBranchOp, MirConstantOp, MirDivOp, MirEqOp, MirFloatConstantOp, MirFuncOp, MirGeOp, MirGotoOp, MirGtOp, MirLeOp, @@ -45,6 +47,9 @@ pub fn register_mir_core_pack(registry: &mut TranslationRegistry) -> Result<(), // 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"))?; @@ -125,6 +130,7 @@ impl OperationTranslation for FunctionTranslation { let mut target = renamed("func.func", input)?; move_property(&mut target, "sym_name", "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)?; if let Some(marker) = target.attributes.remove("gpu_kernel") { match marker { @@ -137,7 +143,7 @@ impl OperationTranslation for FunctionTranslation { .insert("gpu.kernel".into(), MlirAttribute::Unit); move_exact_block_contract(&mut target)?; move_cluster_contract(&mut target)?; - flatten_kernel_slice_abi(ctx, source, &mut target, session)?; + flatten_kernel_slice_abi(ctx, source, &mut target, session, &grid_constants)?; } other => { return Err(format!( @@ -264,14 +270,22 @@ fn take_positive_i32_attribute( Ok(Some(value)) } -/// Match CUDA Oxide's existing kernel launch ABI: a Rust slice is passed as -/// two driver arguments, pointer then length. Rebuild the source aggregate at -/// entry so the already-translated body does not need any rewriting. +/// 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())?; @@ -319,11 +333,13 @@ fn flatten_kernel_slice_abi( let mut flattened_types = Vec::new(); let mut flattened_arguments = Vec::new(); let mut prologue = Vec::new(); + let mut argument_attributes = Vec::new(); - for ((source_type, target_type), original_argument) in source_inputs + 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(); @@ -339,6 +355,10 @@ fn flatten_kernel_slice_abi( 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; } @@ -352,6 +372,10 @@ fn flatten_kernel_slice_abi( 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, @@ -428,6 +452,17 @@ fn flatten_kernel_slice_abi( 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(()) } diff --git a/crates/cuda-oxide-mlir-export/src/mir_memory.rs b/crates/cuda-oxide-mlir-export/src/mir_memory.rs index e61625e966..8996b39210 100644 --- a/crates/cuda-oxide-mlir-export/src/mir_memory.rs +++ b/crates/cuda-oxide-mlir-export/src/mir_memory.rs @@ -704,8 +704,13 @@ impl OperationTranslation for AssertTranslation { _session: &mut TranslationSession<'_>, ) -> Result, String> { ensure_empty_attributes("mir.assert", &input)?; - if input.operands.is_empty() || input.successors.len() != 1 { - return Err("mir.assert expected a condition and one success block".into()); + 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()); } let mut assertion = MlirOperation::new("cf.assert")?; assertion.operands.push(input.operands[0].clone()); @@ -715,11 +720,9 @@ impl OperationTranslation for AssertTranslation { ); assertion.location = input.location.clone(); - let mut branch = MlirOperation::new("cf.br")?; - branch.operands.extend_from_slice(&input.operands[1..]); - branch.successors = input.successors; - branch.location = input.location; - Ok(vec![assertion, branch]) + // Assertions are effectful checks within a block. Any following MIR + // branch is translated separately, so no synthetic branch belongs here. + Ok(vec![assertion]) } } @@ -904,7 +907,7 @@ fn source_is_zero_sized(ctx: &Context, source: TypeHandle) -> bool { false } -fn source_stored_size(ctx: &Context, source: TypeHandle) -> Option { +pub(crate) fn source_stored_size(ctx: &Context, source: TypeHandle) -> Option { let source_ref = source.deref(ctx); if let Some(integer) = source_ref.downcast_ref::() { return Some(u64::from(integer.width()).div_ceil(8)); @@ -1203,7 +1206,7 @@ mod tests { use dialect_mir::{ attributes::{CompilerResultBundleAttr, FieldIndexAttr}, ops::{ - MirAddOp, MirAllocaOp, MirArrayElementAddrOp, MirConstructArrayOp, + MirAddOp, MirAllocaOp, MirArrayElementAddrOp, MirAssertOp, MirConstructArrayOp, MirExtractArrayElementOp, MirExtractFieldOp, MirFieldAddrOp, MirFuncOp, MirLoadOp, MirReturnOp, MirStoreOp, }, @@ -1241,6 +1244,47 @@ mod tests { IntegerType::get(ctx, width, Signedness::Unsigned) } + #[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 inputs = vec![condition.into(), word.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); + 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![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("\"cf.assert\"").count(), 1, "{text}"); + assert!( + text.find("\"cf.assert\"").unwrap() < text.find("\"arith.addi\"").unwrap(), + "{text}" + ); + assert!(!text.contains("\"cf.br\""), "{text}"); + } + #[test] fn compiler_result_array_preserves_lane_order_and_empty_arrays() { for count in [0, 4, 32] { @@ -1250,7 +1294,7 @@ mod tests { 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(&mut ctx, args.clone(), vec![array.into()]); + let ty = FunctionType::get(&ctx, args.clone(), vec![array.into()]); let function_op = Operation::new( &mut ctx, MirFuncOp::get_concrete_op_info(), diff --git a/crates/cuda-oxide-mlir-export/src/nvvm_tcgen05.rs b/crates/cuda-oxide-mlir-export/src/nvvm_tcgen05.rs index ce624d165f..71f3f92e28 100644 --- a/crates/cuda-oxide-mlir-export/src/nvvm_tcgen05.rs +++ b/crates/cuda-oxide-mlir-export/src/nvvm_tcgen05.rs @@ -487,10 +487,10 @@ mod tests { Some("v1:i0763"), ); let mma = Tcgen05MmaOp::new(op); - mma.set_attr_nvvm_tcgen05_mma_form(&mut ctx, Tcgen05MmaFormAttr::Shared); - mma.set_attr_nvvm_tcgen05_mma_kind(&mut ctx, Tcgen05MmaKindAttr::F16); - mma.set_attr_nvvm_tcgen05_mma_cta_group(&mut ctx, Tcgen05MmaCtaGroupAttr::Cg2); - mma.set_attr_nvvm_tcgen05_mma_collector_a(&mut ctx, collector); + 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::( @@ -565,16 +565,16 @@ mod tests { #[test] fn unsupported_mma_forms_fail_instead_of_losing_their_selectors() { - let (mut ctx, module, first) = fixture(); + let (ctx, module, first) = fixture(); let mma = Tcgen05MmaOp::new(first); - mma.set_attr_nvvm_tcgen05_mma_kind(&mut ctx, Tcgen05MmaKindAttr::Tf32); + mma.set_attr_nvvm_tcgen05_mma_kind(&ctx, Tcgen05MmaKindAttr::Tf32); let error = translate(&ctx, &module).unwrap_err(); assert!( error.contains("shared-A FP16 CTA-group-2 MMA only"), "{error}" ); - mma.set_attr_nvvm_tcgen05_mma_kind(&mut ctx, Tcgen05MmaKindAttr::F16); - mma.set_attr_nvvm_tcgen05_mma_cta_group(&mut ctx, Tcgen05MmaCtaGroupAttr::Cg1); + mma.set_attr_nvvm_tcgen05_mma_kind(&ctx, Tcgen05MmaKindAttr::F16); + mma.set_attr_nvvm_tcgen05_mma_cta_group(&ctx, Tcgen05MmaCtaGroupAttr::Cg1); let error = translate(&ctx, &module).unwrap_err(); assert!( error.contains("shared-A FP16 CTA-group-2 MMA only"), diff --git a/crates/cuda-oxide-mlir-export/src/profile.rs b/crates/cuda-oxide-mlir-export/src/profile.rs index f45652fb4d..c3bbe5810b 100644 --- a/crates/cuda-oxide-mlir-export/src/profile.rs +++ b/crates/cuda-oxide-mlir-export/src/profile.rs @@ -79,10 +79,10 @@ pub trait MlirConsumerProfile { } } -/// Full CuTe dialect accepted by the pinned official CUTLASS 4.7 compiler. +/// CuTe mappings supported by the pinned CUTLASS 4.7 compiler. /// -/// This is intentionally versioned. A different public CUTLASS release is a -/// different profile until its syntax, lowering, and runtime gates pass. +/// 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 { gpu_arch: String, @@ -134,8 +134,8 @@ impl MlirConsumerProfile for CutlassFullCuteMlir22 { } fn verify_source_module(&self, ctx: &Context, module: &ModuleOp) -> Result<(), ProfileError> { - // Every backend enters through the same immutable whole-module CuTe - // graph/provenance checks before selecting its continuation. + // 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())) } diff --git a/crates/cuda-oxide-mlir-export/tests/grid_constant.rs b/crates/cuda-oxide-mlir-export/tests/grid_constant.rs new file mode 100644 index 0000000000..e822b6569e --- /dev/null +++ b/crates/cuda-oxide-mlir-export/tests/grid_constant.rs @@ -0,0 +1,190 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use std::num::NonZero; + +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; +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 pliron_mlir_export::{MlirAttribute, render_module}; +use reserved_oxide_symbols::{ + MIR_GRID_CONSTANT_ALIGN_ATTR_PREFIX, MIR_GRID_CONSTANT_POINTEE_ATTR_PREFIX, +}; + +#[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(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), attributes.clone())) +} + +#[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}" + ); +} + +#[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/crates/mir-importer/src/translator/cute/mod.rs b/crates/mir-importer/src/translator/cute/mod.rs index be3908054a..af2ea584a3 100644 --- a/crates/mir-importer/src/translator/cute/mod.rs +++ b/crates/mir-importer/src/translator/cute/mod.rs @@ -3,31 +3,26 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! Recognition of `cute-rs` device-library calls. +//! Translate `cute-rs` calls into CuTe operations. //! -//! ALL cute-specific importer code lives in this module; the shared -//! translator contains only the one dispatch line in `translate_call`. -//! Mirrors the `terminator/intrinsics/` pattern. +//! CuTe-specific call handling lives here. The shared translator calls this +//! module from `translate_call`, as it does for CUDA intrinsics. //! -//! Recognition is STRUCTURAL, not string-based: we rebuild the callee's -//! canonical definition path by walking the DefId parent chain (the -//! cuda-intrinsics mechanism) and compare it against [`CUTE_FNS`]. This is -//! immune to the two failure modes of name matching: +//! 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: //! -//! - rustc's name printer prefers public re-export paths -//! (`cute_rs::load_tile`) over definition paths -//! (`cute_rs::tile::load_tile`), and re-exports are API surface that -//! moves; -//! - trait-method paths print as `::method`, which -//! prefix checks miss entirely. +//! ```text +//! public name: cute_rs::load_tile +//! definition path: cute_rs::tile::load_tile -> CUTE_FNS entry +//! ``` //! -//! Contract: a recognized call is NEVER translated body-by-body. Static -//! parameters are read from the monomorphized substs after validating the -//! per-function [`SubstKind`] schema. Explicit source-abstraction modules -//! contain ordinary device-safe Rust and fall through to body translation; -//! LLVM then erases their `#[inline(always)]` boundaries. Every other unknown -//! `cute_rs` call is a hard error (body translation of a recognition stub -//! would scalarize the layout: the raising trap). +//! A recognized call becomes a CuTe operation; its Rust body is not translated. +//! The importer checks its specialized generic arguments against [`SubstKind`] +//! before reading static configuration. Ordinary Rust helpers in explicitly +//! allowed modules use normal body translation and are inlined by LLVM. +//! Any other unknown `cute_rs` call is an error: translating a compiler stub's +//! body could turn layout operations into scalar code and lose their meaning. pub(crate) mod block_scaled; pub(crate) mod block_scaled_emit; @@ -130,18 +125,15 @@ pub(crate) enum SubstKind { Lifetime, } -/// One recognized function: its canonical DEFINITION path (what the DefId -/// parent walk reconstructs; re-exports never appear here) and the shape -/// its generics must have. +/// A recognized function's definition path and expected generic arguments. +/// Re-export paths do not appear here. pub(crate) struct CuteFnSpec { pub canonical: &'static str, pub func: CuteFn, pub schema: &'static [SubstKind], } -/// The recognition table. Grows by one entry per recognized function; when -/// the op set stabilizes this graduates to catalog generation (the -/// cuda-intrinsics-gen model) and these entries become generated. +/// One entry per recognized function, including its generic argument schema. pub(crate) const CUTE_FNS: &[CuteFnSpec] = &[ CuteFnSpec { canonical: "cute_rs::sm100::__compiler::sm100_tma_store", diff --git a/crates/mir-lower/src/convert/ops/control_flow.rs b/crates/mir-lower/src/convert/ops/control_flow.rs index c086fcb3dc..4e4f899cf4 100644 --- a/crates/mir-lower/src/convert/ops/control_flow.rs +++ b/crates/mir-lower/src/convert/ops/control_flow.rs @@ -322,6 +322,7 @@ mod tests { use crate::convert::ops::test_util::*; use dialect_mir::ops as mir; use dialect_mir::types::{MirPtrType, MirStructType, MirTupleType, StructAbiKind}; + use llvm_export::op_interfaces::VolatilityOpInterface; use llvm_export::ops as llvm; use pliron::builtin::op_interfaces::{ BranchOpInterface, CallOpCallable, CallOpInterface, OperandSegmentInterface, @@ -1056,7 +1057,7 @@ mod tests { assert_eq!(count_ops::(&ctx, &blocks), 3); for (store, expected) in stores.iter().zip([carried, middle, final_value]) { assert_eq!(store.get_operation().deref(&ctx).get_operand(0), expected); - assert!(llvm_export::ops::op_volatile(&ctx, store.get_operation())); + assert!(store.is_volatile(&ctx)); } for (index, condition) in [first_condition, second_condition].into_iter().enumerate() { let block = stores[index] diff --git a/crates/rustc-codegen-cuda/Cargo.lock b/crates/rustc-codegen-cuda/Cargo.lock index 1f593cc64c..dc206d1d27 100644 --- a/crates/rustc-codegen-cuda/Cargo.lock +++ b/crates/rustc-codegen-cuda/Cargo.lock @@ -239,6 +239,7 @@ dependencies = [ "llvm-export", "pliron", "pliron-mlir-export", + "reserved-oxide-symbols", "serde", "serde_json", "thiserror", diff --git a/cuda-oxide-book/compiler/architecture-overview.md b/cuda-oxide-book/compiler/architecture-overview.md index 9b19aa5ae3..242fffc7e7 100644 --- a/cuda-oxide-book/compiler/architecture-overview.md +++ b/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 | +| `pliron-mlir-export` | Typed deterministic Pliron-to-MLIR translation | +| `cuda-oxide-mlir-export` | CUTLASS 4.7 CuTe MLIR mapping packs and profile | | `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 | +| `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` | @@ -150,7 +155,7 @@ cuda-oxide is split into focused crates. Here is every one and its role: | `dialect-ptx` | pliron dialect modelling structured PTX | | `ptx-parse` | Lossless structural views over PTX source text | | `ptx-schedule` | PTX schedule-perturbation fuzzing (nanosleep injection campaigns) | -| `cuda-oxide-codegen` | Rustc-independent PTX backend | +| `cuda-oxide-codegen` | Rustc-independent PTX and CUTLASS cubin backend | | `libnvvm-sys` | `dlopen` bindings to libNVVM (used by `cuda-host::ltoir`) | | `cuda-target-spec` | Shared CUDA target parsing and recorded LLVM PTX-floor policy | | `nvjitlink-sys` | `dlopen` bindings to nvJitLink (used by `cuda-host::ltoir`) | diff --git a/cuteir/README.md b/cuteir/README.md index 326ab48302..b7678014ba 100644 --- a/cuteir/README.md +++ b/cuteir/README.md @@ -1,9 +1,7 @@ # CuTe-style GPU kernels in Rust -`cuteir` lets a Rust kernel describe GPU data by its meaning, not only by its -address. - -The two basic ideas are small: +`cuteir` lets a Rust kernel describe GPU data with typed tensors and layouts. +It starts with two ideas: - a **Tensor** is storage plus a layout; - a **Layout** says where a logical coordinate lives in that storage. @@ -40,12 +38,14 @@ Tensor → zipped_divide → slice → copy / add / matrix multiply preparation + verification │ ▼ - selected backend continuation + code generation ``` -The whole-module verifier checks that tensor provenance, scheduler and -pipeline state, TMA transactions, shared-memory MMA, and epilogue operations -form a complete semantic story. It does not choose or execute a backend. +The verifier checks where tensors come from, how scheduler and pipeline state +flows between operations, and whether TMA, MMA, and epilogue operations agree +on layouts and buffer ownership. Both backends use these checks before +lowering CuTe operations. Dynamic phases, lane participation, and pointer +lifetimes remain the caller's responsibility. ## What lives here @@ -85,13 +85,10 @@ The elementwise example uses ordinary `f32` and `f16` operations. The two FP4 examples use Blackwell instructions and should be built for `sm_120a`; their own READMEs give the exact device and shape requirements. -The first three programs check every GPU output bit pattern against a deterministic host -result before reporting success. The FP16 example checks every output -against a numerical oracle; its README includes the SM100 build and validation commands. -Backend-specific artifact and code-shape -checks live with the backend implementation rather than in this shared layer. +The first three programs compare every GPU output bit pattern with a +fixed host result. The FP16 example checks outputs numerically; its README +includes the SM100 build and validation commands. Tests for generated code +and artifacts live with each backend. -On the translation branch, the [official CUTLASS backend](docs/translation.md) -maps these semantic operations directly to the CUTLASS 4.7 MLIR profile and -embeds the compiler-produced cubin. It does not pass through native CuTe -expansion. +The [official CUTLASS backend](docs/translation.md) maps CuTe operations +directly to the CUTLASS 4.7 MLIR profile and embeds the resulting cubin. diff --git a/cuteir/cute-rs/src/sm100.rs b/cuteir/cute-rs/src/sm100.rs index b21d26a19e..662deaca7b 100644 --- a/cuteir/cute-rs/src/sm100.rs +++ b/cuteir/cute-rs/src/sm100.rs @@ -3,11 +3,15 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! SM100 two-CTA tensor-memory GEMM building blocks. +//! SM100 GEMM building blocks for a pair of thread blocks (CTAs). //! -//! Shared layouts, the two N partitions, and the A collector lifetime are -//! compiler-visible semantics. The scalar compiler boundary below is lowered -//! to dialect-cute and then to CUTLASS; it has no raw CUDA implementation. +//! 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 +//! ``` use crate::TmaDesc; use crate::markers::ReifySmem2D; @@ -29,6 +33,7 @@ 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. @@ -52,6 +57,7 @@ impl Sm100Tmem { /// 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. @@ -64,6 +70,7 @@ impl Sm100Tmem { /// 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)] @@ -75,9 +82,15 @@ impl Sm100Tmem { /// Two-CTA FP16 tiled MMA with two N partitions sharing the A collector. /// -/// Each K=16 step fills A with the first N partition and consumes it for -/// the last time with the second. The collector lifetime never crosses a K -/// step. The complete output is `256 x (N0 + N1)` with FP32 accumulation. +/// 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. pub struct Sm100TiledMma( PhantomData<(AL, B0L, B1L)>, ); @@ -91,15 +104,15 @@ impl< > Sm100TiledMma { const VALID: () = { - assert!(N0 >= 16 && N0 <= 256 && N0 % 16 == 0); - assert!(N1 >= 16 && N1 <= 256 && N1 % 16 == 0); + 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::ROWS == 128 && AL::COLS == K as i64); assert!(B0L::ROWS == (N0 / 2) as i64 && B0L::COLS == K as i64); assert!(B1L::ROWS == (N1 / 2) as i64 && B1L::COLS == K as i64); }; - /// Define the compile-time tile and collector contract. + /// Check the tile dimensions and create the MMA configuration. #[inline(always)] pub const fn new() -> Self { let () = Self::VALID; @@ -109,6 +122,7 @@ impl< /// 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 @@ -152,12 +166,13 @@ impl Sm100ClusterTmaCopy { /// 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 - /// its issuing lane internally and may synchronize - /// the warp. Calling it from an already elected lane can deadlock when + /// 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 @@ -177,12 +192,13 @@ impl Sm100ClusterTmaCopy { } } -/// Drain one 128x32 FP32 accumulator slice to a laid-out FP16 shared tile. +/// Convert one 128x32 FP32 accumulator slice to an FP16 tile in shared memory. pub struct Sm100TmemEpilogue(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 @@ -204,6 +220,7 @@ impl Sm100TmemEpilogue { /// 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 @@ -223,8 +240,9 @@ impl Sm100TmemEpilogue { /// Two-CTA TMA-to-MMA stage ring, with one completion arrival per stage. /// -/// `TX_BYTES` includes every copy from both CTAs. The elected leader produces -/// the expectation, and asynchronous MMA completion releases both CTAs. +/// `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 Sm100TmaMmaPipeline { full: *mut Barrier, empty: *mut Barrier, @@ -234,6 +252,7 @@ impl Sm100TmaMmaPipeline Sm100TmaMmaPipeline Sm100TmaMmaPipeline Sm100TmaMmaPipeline Sm100TmaMmaPipeline Sm100TmaMmaPipeline Sm100TmaMmaPipeline Sm100AccumulatorPipeline { /// 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)] @@ -312,6 +338,7 @@ impl Sm100AccumulatorPipeline { /// 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) { @@ -320,6 +347,7 @@ impl Sm100AccumulatorPipeline { /// 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) { @@ -328,6 +356,7 @@ impl Sm100AccumulatorPipeline { /// 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) { @@ -336,6 +365,7 @@ impl Sm100AccumulatorPipeline { /// 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) { @@ -344,6 +374,7 @@ impl Sm100AccumulatorPipeline { /// 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. @@ -381,7 +412,7 @@ impl Sm100TmaStorePipeline { } } -/// Stable semantic compiler boundaries; never execute on the host. +/// Calls recognized by the compiler; never execute these on the host. #[doc(hidden)] #[allow(missing_docs)] pub mod __compiler { @@ -412,8 +443,8 @@ pub mod __compiler { let _ = (tmem, dst, tid, bias, PhantomData::); unreachable!("SM100 CuTe marker requires device compilation") } - // Collective boundary: every lane of one producer warp participates. - // The native CuTe load owns election and any required warp synchronization. + // Every lane of the producer warp participates. CuTe selects the issuing + // lane and performs any required warp synchronization. #[inline(never)] pub unsafe fn sm100_cluster_tma_load( desc: *const TmaDesc, diff --git a/cuteir/cute-rs/src/tma.rs b/cuteir/cute-rs/src/tma.rs index 49d248700d..823560c31d 100644 --- a/cuteir/cute-rs/src/tma.rs +++ b/cuteir/cute-rs/src/tma.rs @@ -15,9 +15,9 @@ //! the descriptor type: //! //! ```text -//! host: let desc: TmaDesc = make_tma_desc_2d(...)?; -//! kernel: fn kernel(desc: *const TmaDesc, ...) -//! └── same SmemL is required +//! host: TmaDesc -> kernel: &TmaDesc +//! descriptor value #[grid_constant] parameter +//! same element and layout types //! ``` //! //! If host and kernel layouts differ, the program does not compile. A @@ -47,6 +47,10 @@ 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: &TmaDesc` 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 TmaDesc { /// Descriptor bytes produced by the CUDA driver. @@ -55,6 +59,15 @@ pub struct TmaDesc { pub marker: PhantomData<(T, SmemL)>, } +// The type parameters describe the encoded bytes; they are not stored values. +impl Copy for TmaDesc {} + +impl Clone for TmaDesc { + fn clone(&self) -> Self { + *self + } +} + /// Shared-memory swizzles supported by TMA. /// /// TMA supports only these byte-based patterns. Each keeps 16-byte chunks @@ -81,9 +94,8 @@ pub enum TmaSwizzleMode { /// Requested level-2 (L2) cache fetch size for a TMA descriptor. /// -/// The default is [`Self::None`] to preserve the original behavior of -/// `make_tma_desc_2d`. CUTLASS and CuTeDSL commonly request [`Self::B128`] -/// for GEMM input tiles. +/// [`Self::None`] is the default used by `make_tma_desc_2d`. 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. @@ -104,12 +116,12 @@ pub struct TmaEncodeOptions { } impl TmaEncodeOptions { - /// Keep the original policy: do not request a larger L2 fetch. + /// Do not request a larger L2 fetch. pub const DEFAULT: Self = Self { l2_promotion: TmaL2Promotion::None, }; - /// Match the usual CUTLASS/CuTeDSL input policy: fetch 128-byte sectors. + /// Request 128-byte L2 sectors. pub const L2_128B: Self = Self { l2_promotion: TmaL2Promotion::B128, }; @@ -440,6 +452,16 @@ pub use host::{TmaElement, make_tma_desc_2d, make_tma_desc_2d_with_options}; #[cfg(test)] mod tests { + #[test] + fn descriptor_value_abi_does_not_require_copy_marker_types() { + // Layout and element markers need not be Copy: only descriptor bytes move. + struct Marker; + fn assert_copy() {} + assert_copy::>(); + assert_eq!(core::mem::size_of::>(), 128); + assert_eq!(core::mem::align_of::>(), 64); + } + use super::*; #[test] diff --git a/cuteir/dialect-cute/src/lib.rs b/cuteir/dialect-cute/src/lib.rs index 55775320d2..edf6ea76c8 100644 --- a/cuteir/dialect-cute/src/lib.rs +++ b/cuteir/dialect-cute/src/lib.rs @@ -5,20 +5,19 @@ //! CuTe dialect definition. //! -//! Models the CuTe layout algebra as first-class pliron ops, types, and -//! attributes, coexisting in one module with dialect-mir control flow and -//! scalar operations. +//! 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); result types of algebra ops are computed with `cute-layout` at -//! IR-build time. The backend-neutral verifier checks the complete semantic -//! graph before any backend continuation begins. +//! 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. -//! These semantic operations are consumed by the selected backend continuation. +//! The selected backend lowers these operations to GPU code. pub mod attributes; pub mod epilogue_ops; @@ -33,9 +32,7 @@ pub mod tensor_ops; pub mod types; pub mod verify; -/// Re-export of the layout algebra so consumers (mir-importer's recognition -/// module and backend continuations) depend only on `dialect-cute` and reach -/// the algebra through it. +/// Layout math shared by the importer, verifier, and backends. pub use cute_layout as layout; use pliron::attribute::Attribute; @@ -55,7 +52,7 @@ pub fn register(ctx: &mut Context) { sm100_ops::register(ctx); // The #[pliron_op]/#[pliron_attr] macros auto-register via Context::default; - // explicit registration is the idempotent house convention. + // registering them again is safe and matches the other dialects. ops::CuteCopyOp::register(ctx); ops::CuteAssumeDivOp::register(ctx); ops::CuteCopyG2SOp::register(ctx); diff --git a/cuteir/dialect-cute/src/sm100_ops.rs b/cuteir/dialect-cute/src/sm100_ops.rs index 185f0b7b10..8da194842a 100644 --- a/cuteir/dialect-cute/src/sm100_ops.rs +++ b/cuteir/dialect-cute/src/sm100_ops.rs @@ -115,8 +115,8 @@ impl Verify for CuteSm100MmaPlanAttr { || self.tmem_columns != 512 || !(16..=256).contains(&self.n0) || !(16..=256).contains(&self.n1) - || self.n0 % 16 != 0 - || self.n1 % 16 != 0 + || !self.n0.is_multiple_of(16) + || !self.n1.is_multiple_of(16) || self.n0 + self.n1 > self.tmem_columns { return verify_err_noloc!( @@ -218,13 +218,13 @@ impl Verify for CuteSm100PipelinePlanAttr { if self.stages != 1 || self.consumer_arrivals == 0 || self.consumer_arrivals > 64 - || self.consumer_arrivals % 2 != 0 + || !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 % 16 != 0 { + } 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" ); diff --git a/cuteir/dialect-cute/src/types.rs b/cuteir/dialect-cute/src/types.rs index 53b589dfaf..94f5e946e5 100644 --- a/cuteir/dialect-cute/src/types.rs +++ b/cuteir/dialect-cute/src/types.rs @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! Tensor-view types kept only while CuTe meaning is useful. +//! Tensor-view types used to check and lower CuTe operations. //! //! A tensor view is a compiler description, not a runtime struct: //! @@ -14,20 +14,19 @@ //! ``` //! //! A selected backend consumes the description while lowering to native -//! operations or MLIR. No backend may assign this ghost type an ABI or leave -//! it in a kernel signature. +//! operations or MLIR. A view has no runtime storage and must be removed +//! before the backend produces a kernel signature. //! -//! Backend-neutral verification reads these values after common MIR -//! preparation has exposed the semantic story: +//! After MIR preparation, the verifier follows each view through its uses: //! //! ```text //! make -> zipped_divide -> slice -> load/store //! ``` //! -//! At that seam, a view cannot sit inside a pointer or aggregate, cross a -//! block edge, or appear in a function signature. Ordinary scalar carriers -//! may remain in closed compiler-owned local cells or cross CFG edges; the -//! semantic verifier follows those carriers before backend selection. +//! 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. use dialect_mir::types::{MirArrayType, MirFP16Type, MirStructType, MirTupleType}; use pliron::builtin::types::{FP32Type, IntegerType}; @@ -47,9 +46,8 @@ use crate::attributes::{ /// Physical scalar registers inside one ordinary MIR aggregate. /// -/// Semantic MMA operations keep their runtime values in the same Rust-shaped -/// aggregates used today. This small summary lets their verifiers check those -/// carriers without turning them into new loop-carried ghost values. +/// 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 { pub u32_registers: u64, @@ -101,7 +99,7 @@ impl MmaCarrierShape { /// /// Arrays, tuples, and structs are walked recursively. Empty marker structs /// such as `PhantomData` contribute nothing. Any other scalar or -/// aggregate kind fails closed with `None`. +/// aggregate kind returns `None`. #[must_use] pub fn mma_carrier_shape(ctx: &Context, ty: TypeHandle) -> Option { let ty = ty.deref(ctx); diff --git a/cuteir/dialect-cute/src/verify.rs b/cuteir/dialect-cute/src/verify.rs index f5e8162c0a..577231e8e9 100644 --- a/cuteir/dialect-cute/src/verify.rs +++ b/cuteir/dialect-cute/src/verify.rs @@ -3,11 +3,11 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! Backend-neutral verification for complete high-level CuTe stories. +//! Check CuTe values and operation order before either backend lowers them. //! -//! This module is deliberately analysis-only. It follows semantic SSA -//! provenance, checks closed use graphs and ordered protocols, and audits -//! compiler-only CuTe types without cloning or rewriting the input module. +//! 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. use std::collections::{HashMap, HashSet}; diff --git a/cuteir/docs/translation.md b/cuteir/docs/translation.md index c54ad920e7..c767f18de2 100644 --- a/cuteir/docs/translation.md +++ b/cuteir/docs/translation.md @@ -1,11 +1,11 @@ # Official CUTLASS translation backend -Translate Rust CuTe contracts into NVIDIA's pinned CUTLASS 4.7 MLIR profile: +Compile Rust CuTe kernels with NVIDIA's pinned CUTLASS 4.7 compiler: ![CuTe compiler flow from Rust layouts through verified MLIR and CUTLASS to a CUDA cubin](assets/cutlass-translation-flow.svg) -- High-level CuTe survives preparation; reviewed mapping packs translate it - directly, without native CuTe expansion or prior MIR/NVVM leaf lowering. +- Shared preparation keeps CuTe operations intact. The mapping packs translate + them directly to the MLIR accepted by CUTLASS. - Ordinary non-CuTe kernels retain the default MIR/NVVM/LLVM-to-PTX backend. ## Install and build @@ -42,8 +42,23 @@ CUDA_OXIDE_CUTLASS_COMPILER=/absolute/path/to/libCutlassCompiler.so \ | [Block-scaled GEMM](../examples/blockscale_gemm_cute) | Scheduler, work tiles, TMA pipelines, shared-memory MMA, epilogue stores | | [SM100 FP16 GEMM](../examples/fp16_gemm_256x352_cute) | Typed shared tiles, TMEM, paired two-CTA MMA, cluster TMA, producer/consumer pipelines | -Block-scaled epilogue hand-offs: `ReadyForTma` emits an async-shared proxy -publication fence, then a counted CTA barrier; `Reusable` emits only the barrier. +At a block-scaled epilogue hand-off, `ReadyForTma` makes shared-memory writes +visible to TMA with an async-shared proxy fence, then emits a counted CTA +barrier. `Reusable` emits only the barrier. + +Grid-constant references use the same launch convention as the native backend: + +```text +#[grid_constant] desc: &TmaDesc<...> +host descriptor -> 128-byte launch argument -> device &TmaDesc<...> + aligned to 64 bytes same address for every thread +``` + +The backend checks the cubin's argument sizes and offsets against this layout. +The descriptor needs no separate device allocation; the tensor allocations it +points to must remain valid until GPU work completes. See the +[grid-constant safety requirements](../../cuda-oxide-book/gpu-programming/kernels-and-device-functions.md#grid-constant-parameters) +for the full launch contract. ## SM100 contracts @@ -56,19 +71,20 @@ The [Rust APIs](../cute-rs/src/sm100.rs) become verified `cute.sm100_*` plans: | `Sm100TmaMmaPipeline`, `Sm100AccumulatorPipeline` | Buffer ownership and completion | | `Sm100Tmem`, `Sm100TmemEpilogue` | TMEM allocation and FP32 → FP16 epilogue | -- **Verified:** layout/partition compatibility, cluster transaction bytes, - barrier arrivals, and allocation/MMA/copy/release lifecycle presence. -- **Caller-owned unsafe preconditions:** dynamic phases, lane participation, - and pointer lifetimes. +- **Compiler checks:** matching layouts and partitions, cluster transaction + bytes, barrier arrivals, and the presence of allocation, MMA, copy, and + release operations. +- **Caller responsibilities:** correct dynamic phases, participating lanes, + and pointer lifetimes at each unsafe operation. - **Profile:** two CTAs; FP16 inputs, FP32 accumulation; M256/K64; 128-byte input swizzle; 512 TMEM columns; 128×32 epilogue with 64-byte swizzle. The example uses N192+160 and checks CPU FP32 accumulation rounded to FP16. -### CuTe and the collector leaf +### CuTe and the A collector -- Native CuTe builds shared descriptors, TMEM operations, and cluster TMA. -- CUTLASS 4.7's native MMA atom lacks A-collector fill/last-use control. - The backend therefore emits an **NVVM TCGen05 MMA leaf** from the verified plan: +- CUTLASS CuTe operations build shared descriptors, TMEM operations, and cluster TMA. +- CUTLASS 4.7's MMA atom cannot select when to fill or release the A collector. + The backend emits an **NVVM TCGen05 MMA instruction** from the verified plan: ```text Each K16: first N partition → fill A collector diff --git a/cuteir/examples/blockscale_gemm_cute/Cargo.lock b/cuteir/examples/blockscale_gemm_cute/Cargo.lock index cd0219bd05..5196c6b64b 100644 --- a/cuteir/examples/blockscale_gemm_cute/Cargo.lock +++ b/cuteir/examples/blockscale_gemm_cute/Cargo.lock @@ -19,14 +19,16 @@ checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" [[package]] name = "bindgen" -version = "0.71.1" +version = "0.69.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" dependencies = [ "bitflags", "cexpr", "clang-sys", "itertools", + "lazy_static", + "lazycell", "log", "prettyplease", "proc-macro2", @@ -35,6 +37,7 @@ dependencies = [ "rustc-hash", "shlex", "syn 2.0.119", + "which", ] [[package]] @@ -135,22 +138,42 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.2.1" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f42e2c99c56bf2eef9d569435bfd72a478dfaea3f178b8e5bae891875bbbe2e" dependencies = [ "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] name = "cuda-core" -version = "0.2.1" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84bdb592a2844148fb13bdfbbab158879f56916bde1d1a1b6580ca27f23850cd" dependencies = [ "anyhow", "cuda-bindings", - "cuda-macros", + "cuda-core-derive", "half", "oxide-artifacts", ] +[[package]] +name = "cuda-core-derive" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a48da7d10c1078f7fbb87d7176e9476affadea148babf31946e348073872d44" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "cuda-device" version = "0.2.1" @@ -220,6 +243,16 @@ 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" @@ -268,6 +301,15 @@ 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" @@ -280,13 +322,25 @@ dependencies = [ [[package]] name = "itertools" -version = "0.13.0" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" dependencies = [ "either", ] +[[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.189" @@ -312,6 +366,12 @@ dependencies = [ "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.33" @@ -360,9 +420,17 @@ dependencies = [ "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", ] @@ -434,9 +502,22 @@ version = "0.2.1" [[package]] name = "rustc-hash" -version = "2.1.3" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" +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" @@ -545,12 +626,106 @@ 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.56" diff --git a/cuteir/examples/blockscale_gemm_cute/Cargo.toml b/cuteir/examples/blockscale_gemm_cute/Cargo.toml index 5a1aa09166..bdc985bde5 100644 --- a/cuteir/examples/blockscale_gemm_cute/Cargo.toml +++ b/cuteir/examples/blockscale_gemm_cute/Cargo.toml @@ -13,4 +13,4 @@ description = "Block-scaled MXFP4 GEMM composing typed TMA pipelines, shared-mem cute-rs = { path = "../../cute-rs", features = ["host"] } cuda-device = { path = "../../../crates/cuda-device" } cuda-host = { path = "../../../crates/cuda-host" } -cuda-core = { path = "../../../crates/cuda-core" } +cuda-core = { version = "0.3.1", features = ["f16"] } diff --git a/cuteir/examples/elementwise_cute/Cargo.lock b/cuteir/examples/elementwise_cute/Cargo.lock index 3d7c4c6c21..f30d5e565a 100644 --- a/cuteir/examples/elementwise_cute/Cargo.lock +++ b/cuteir/examples/elementwise_cute/Cargo.lock @@ -19,14 +19,16 @@ checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" [[package]] name = "bindgen" -version = "0.71.1" +version = "0.69.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" dependencies = [ "bitflags", "cexpr", "clang-sys", "itertools", + "lazy_static", + "lazycell", "log", "prettyplease", "proc-macro2", @@ -35,6 +37,7 @@ dependencies = [ "rustc-hash", "shlex", "syn 2.0.119", + "which", ] [[package]] @@ -125,22 +128,42 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.2.1" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f42e2c99c56bf2eef9d569435bfd72a478dfaea3f178b8e5bae891875bbbe2e" dependencies = [ "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] name = "cuda-core" -version = "0.2.1" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84bdb592a2844148fb13bdfbbab158879f56916bde1d1a1b6580ca27f23850cd" dependencies = [ "anyhow", "cuda-bindings", - "cuda-macros", + "cuda-core-derive", "half", "oxide-artifacts", ] +[[package]] +name = "cuda-core-derive" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a48da7d10c1078f7fbb87d7176e9476affadea148babf31946e348073872d44" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "cuda-device" version = "0.2.1" @@ -219,6 +242,16 @@ 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" @@ -267,6 +300,15 @@ 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" @@ -279,13 +321,25 @@ dependencies = [ [[package]] name = "itertools" -version = "0.13.0" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" dependencies = [ "either", ] +[[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.189" @@ -311,6 +365,12 @@ dependencies = [ "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.33" @@ -359,9 +419,17 @@ dependencies = [ "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", ] @@ -433,9 +501,22 @@ version = "0.2.1" [[package]] name = "rustc-hash" -version = "2.1.3" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" +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" @@ -544,12 +625,106 @@ 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.56" diff --git a/cuteir/examples/elementwise_cute/Cargo.toml b/cuteir/examples/elementwise_cute/Cargo.toml index 3abd0aee2a..91625cfdc6 100644 --- a/cuteir/examples/elementwise_cute/Cargo.toml +++ b/cuteir/examples/elementwise_cute/Cargo.toml @@ -13,4 +13,4 @@ description = "CuTe-style f32 and f16 elementwise add with vectorized full tiles cute-rs = { path = "../../cute-rs" } cuda-device = { path = "../../../crates/cuda-device" } cuda-host = { path = "../../../crates/cuda-host" } -cuda-core = { path = "../../../crates/cuda-core" } +cuda-core = { version = "0.3.1", features = ["f16"] } diff --git a/cuteir/examples/elementwise_cute/src/main.rs b/cuteir/examples/elementwise_cute/src/main.rs index 038d5d9c3a..aad1ae75e0 100644 --- a/cuteir/examples/elementwise_cute/src/main.rs +++ b/cuteir/examples/elementwise_cute/src/main.rs @@ -20,7 +20,8 @@ #![feature(f16)] -use cuda_core::{CudaContext, DeviceBuffer, LaunchConfig}; +use cuda_core::simt::LaunchConfig; +use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; /// Deliberately not divisible by the tiles: f32 tail of 3, f16 tail of 5. diff --git a/cuteir/examples/fp16_gemm_256x352_cute/.gitignore b/cuteir/examples/fp16_gemm_256x352_cute/.gitignore index 5613b74e05..b49f80c109 100644 --- a/cuteir/examples/fp16_gemm_256x352_cute/.gitignore +++ b/cuteir/examples/fp16_gemm_256x352_cute/.gitignore @@ -1,6 +1,7 @@ # Local Python reference and comparison harness. /reference/ /compare.py +/.venv # Local performance investigations and benchmark artifacts. /benchmark-results/ diff --git a/cuteir/examples/fp16_gemm_256x352_cute/Cargo.lock b/cuteir/examples/fp16_gemm_256x352_cute/Cargo.lock index e55d8136c7..fc19c6c6e3 100644 --- a/cuteir/examples/fp16_gemm_256x352_cute/Cargo.lock +++ b/cuteir/examples/fp16_gemm_256x352_cute/Cargo.lock @@ -19,14 +19,16 @@ checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" [[package]] name = "bindgen" -version = "0.71.1" +version = "0.69.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" dependencies = [ "bitflags", "cexpr", "clang-sys", "itertools", + "lazy_static", + "lazycell", "log", "prettyplease", "proc-macro2", @@ -35,6 +37,7 @@ dependencies = [ "rustc-hash", "shlex", "syn 2.0.119", + "which", ] [[package]] @@ -125,22 +128,42 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.2.1" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f42e2c99c56bf2eef9d569435bfd72a478dfaea3f178b8e5bae891875bbbe2e" dependencies = [ "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] name = "cuda-core" -version = "0.2.1" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84bdb592a2844148fb13bdfbbab158879f56916bde1d1a1b6580ca27f23850cd" dependencies = [ "anyhow", "cuda-bindings", - "cuda-macros", + "cuda-core-derive", "half", "oxide-artifacts", ] +[[package]] +name = "cuda-core-derive" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a48da7d10c1078f7fbb87d7176e9476affadea148babf31946e348073872d44" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "cuda-device" version = "0.2.1" @@ -210,6 +233,16 @@ 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" @@ -269,6 +302,15 @@ 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" @@ -281,9 +323,9 @@ dependencies = [ [[package]] name = "itertools" -version = "0.13.0" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" dependencies = [ "either", ] @@ -294,6 +336,18 @@ 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.189" @@ -319,6 +373,12 @@ dependencies = [ "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" @@ -367,9 +427,17 @@ dependencies = [ "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", ] @@ -441,9 +509,22 @@ version = "0.2.1" [[package]] name = "rustc-hash" -version = "2.1.3" +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 = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] [[package]] name = "serde" @@ -565,12 +646,106 @@ 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.57" diff --git a/cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml b/cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml index 34adaaef83..2b8799763e 100644 --- a/cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml +++ b/cuteir/examples/fp16_gemm_256x352_cute/Cargo.toml @@ -11,5 +11,5 @@ description = "SM100 two-CTA persistent FP16 GEMM with a 256x352 output tile" cute-rs = { path = "../../cute-rs", features = ["host"] } cuda-device = { path = "../../../crates/cuda-device" } cuda-host = { path = "../../../crates/cuda-host" } -cuda-core = { path = "../../../crates/cuda-core" } +cuda-core = "0.3.1" serde_json = "1" diff --git a/cuteir/examples/fp16_gemm_256x352_cute/README.md b/cuteir/examples/fp16_gemm_256x352_cute/README.md index dc2dc2441f..5cf3d556da 100644 --- a/cuteir/examples/fp16_gemm_256x352_cute/README.md +++ b/cuteir/examples/fp16_gemm_256x352_cute/README.md @@ -1,7 +1,8 @@ # FP16 GEMM with CuTe on SM100 -A persistent, two-CTA GEMM expressed with Rust layouts, tiled copies/MMA, -tensor memory, and asynchronous pipelines. +Two thread blocks (CTAs) work together on each 256 × 352 output tile. The +kernel uses Rust types to describe tile layouts, TMA copies, and matrix +multiply-and-accumulate (MMA) operations. ```text C[M, N] = A[M, K] @ B[N, K].T + bias[M, 1] # bias is optional @@ -9,7 +10,10 @@ C[M, N] = A[M, K] @ B[N, K].T + bias[M, 1] # bias is optional - A, B, C, and bias are FP16; multiplication accumulates in FP32. - `--has-bias` adds one value per output row before conversion to FP16. -- Layouts and pipeline contracts stay visible through the compiler: +- All four typed TMA descriptors use `#[grid_constant]` parameters. The host + passes their 128-byte values directly, without allocating descriptor buffers + in device global memory. +- The compiler keeps the layout and pipeline information when producing GPU code: ![CuTe compiler flow from Rust layouts through verified MLIR and CUTLASS to a CUDA cubin](../../docs/assets/cutlass-translation-flow.svg) @@ -19,9 +23,9 @@ C[M, N] = A[M, K] @ B[N, K].T + bias[M, 1] # bias is optional one 256 × 352 output tile 192 columns 160 columns ┌─────────────────┬────────────────┐ -CTA 0: 128 rows│ │ │ +CTA 0: 128 rows │ │ │ ├─────────────────┼────────────────┤ -CTA 1: 128 rows│ │ │ +CTA 1: 128 rows │ │ │ └─────────────────┴────────────────┘ two-CTA MMA pair; A collector reused @@ -52,33 +56,83 @@ C global ◀─────TMA────── shared [2 stages] ◀──FP ## Build and run Requirements: an SM100 GPU, a compatible CUDA driver, and the repository's -configured cuda-oxide toolchain. From the repository root: +configured cuda-oxide toolchain. The default comparison also needs the local +CuTe DSL assets and Python environment described below. From the repository root: ```bash cargo oxide toolchain install cutlass export CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir cargo oxide build fp16_gemm_256x352_cute --arch sm_100a -# Exactly one kernel launch; verify its output afterward. +# Verify and compare CuTe DSL and cuda-oxide; print two performance lines. +cargo oxide run fp16_gemm_256x352_cute --arch sm_100a + +# Rust only: exactly one kernel launch; verify its output afterward. cargo oxide run fp16_gemm_256x352_cute --arch sm_100a -- \ - --mnk 256,352,64 --timing-mode direct --warmup 0 --iters 1 + --rust-only --mnk 256,352,64 --timing-mode direct --warmup 0 --iters 1 -# N/K tails, row bias, persistent tiles, and five-stage ring reuse. +# Rust only: N/K tails, row bias, persistent tiles, and five-stage ring reuse. cargo oxide run fp16_gemm_256x352_cute --arch sm_100a -- \ - --mnk 512,712,392 --has-bias --clusters 1 \ + --rust-only --mnk 512,712,392 --has-bias --clusters 1 \ --timing-mode direct --warmup 0 --iters 1 ``` - M must be a positive multiple of **256**; N and K, of **8**. TMA handles partial N352 and K64 tiles. -- The runner checks every output against CPU FP32 accumulation rounded to - FP16 (`atol=0.1`, `rtol=1e-5`). Large shapes make this check expensive. -- Standalone timing defaults to **graph**. The commands above explicitly use - **direct**: one kernel per CUDA-event pair, with no graph capture. -- `--warmup 0 --iters 1` in direct mode launches once total and verifies that - timed output. Setup and verification are outside the event interval. -- `--json result.json` saves configuration, verification, and timing metadata; - `--help` lists all options. Graph mode uses `--graph-launches` launches per sample. +- The default comparison uses `256,352,64`, 25 warmup launches, and 100 timed + samples per implementation. Override these with `--mnk`, `--warmup`, and `--iters`. +- `--rust-only` checks every output against CPU FP32 accumulation rounded to + FP16 (`atol=0.1`, `rtol=1e-5`) and prints its JSON report. Large shapes make + this CPU check expensive. +- Rust-only timing defaults to **graph**, with `--graph-launches` launches per + sample. The standalone commands above explicitly use **direct**: one kernel + per CUDA-event pair, with no graph capture. +- `--warmup 0 --iters 1` in direct mode launches each selected implementation + once total and verifies that timed output. Setup and verification are outside + the event interval. +- `--json result.json` saves the combined comparison report by default, or the + standalone cuda-oxide report with `--rust-only`. `--help` lists all options. + +## Compare with CuTe DSL + +Normal execution invokes the separately supplied local `compare.py` and +`reference/` assets; these remain ignored and are not included in a clean +checkout. The Python environment needs `nvidia-cutlass-dsl[cu13]==4.7.0`, +CUDA-enabled PyTorch, CUDA Python, and NumPy. + +The executable selects Python in this order: `--python PATH`, +`CUDA_OXIDE_CUTE_PYTHON`, this example's `.venv/bin/python`, then `python3` on +`PATH`. Placing the environment in `cuteir/examples/fp16_gemm_256x352_cute/.venv` +lets the default command find it automatically. To select another environment +and save results for a larger shape: + +```bash +cargo oxide run fp16_gemm_256x352_cute --arch sm_100a -- \ + --python /path/to/venv/bin/python \ + --mnk 8192,8096,8192 --warmup 100 --iters 31 --json /tmp/gemm-comparison.json +``` + +The comparison verifies both outputs on identical inputs and prints two lines: + +```text +CuTe DSL:
];`. +/// +/// # 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 db8f9f835d..be541602cf 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/movmatrix.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 24409bc43b..7b7d77e5cf 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/prmt.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 8ec46c6814..ba53dbe2d2 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/register_control.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/register_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 6ba412820b..01b82a4570 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/register_mma.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/register_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::wmma` to keep public paths stable. 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 04ed9a6f56..5549a4f3b6 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/shared_sreg.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/shared_sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 06489ee489..9124403944 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/sparse_mma.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::wmma` to keep public paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/sreg.rs b/cuda-oxide/crates/cuda-device/src/generated/sreg.rs index 05589572c3..1cb660d202 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/sreg.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // This file is included lexically inside `cuda_device::thread` so existing diff --git a/cuda-oxide/crates/cuda-device/src/generated/stmatrix.rs b/cuda-oxide/crates/cuda-device/src/generated/stmatrix.rs index c8dbbc3c88..b7f561498c 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/stmatrix.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 51db6d5fd7..92cce88185 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/tcgen05.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/tcgen05.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 ac179e4134..24f33db965 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/tma.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::tma` to keep its public API stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/warp_sreg.rs b/cuda-oxide/crates/cuda-device/src/generated/warp_sreg.rs index 91ebfb089d..a83af970b9 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/warp_sreg.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/warp_sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 2ed7be3489..fb9cc4912b 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/wgmma_control.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::wgmma` to keep its public API stable. diff --git a/cuda-oxide/crates/cuda-host/src/lib.rs b/cuda-oxide/crates/cuda-host/src/lib.rs index e3dcdaddf6..1fa90924a7 100644 --- a/cuda-oxide/crates/cuda-host/src/lib.rs +++ b/cuda-oxide/crates/cuda-host/src/lib.rs @@ -15,7 +15,6 @@ // 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. #![feature(core_intrinsics)] -#![feature(f16)] #![allow(internal_features)] //! Host-side utilities for CUDA kernel development. diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/memory.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/memory.rs index 78d47948ab..7c6edea917 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/memory.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/memory.rs @@ -12,11 +12,14 @@ pub enum MbarrierExtendedOperation { ArriveExpectTxCta, ArriveExpectTxCluster, ArriveRemoteCluster, + ArriveRemoteClusterRelaxed, TryWaitTokenCta, TryWaitParityCta, TryWaitParityCluster, + TryWaitParityClusterRelaxed, FenceProxyAsyncSharedCta, FenceMbarrierInitReleaseCluster, + FenceAcquireSharedClusterCluster, FenceProxyAsyncGenericReleaseSharedCtaCluster, FenceProxyAsyncGenericAcquireSharedClusterCluster, Nanosleep, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/overlay.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/overlay.rs index 081f4b33b7..7bbe0a22f7 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/overlay.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/overlay.rs @@ -166,6 +166,8 @@ pub struct StmatrixAdmissionVariant { pub struct MbarrierExtendedAdmission { pub llvm_evidence_profile: String, pub libnvvm_evidence_profile: String, + #[serde(default)] + pub relaxed_libnvvm_evidence_profile: Option, pub runtime_validation: RuntimeValidation, #[serde(rename = "variant")] pub variants: Vec, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/compat.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/compat.rs index 49239b1e8b..dcdba3aa09 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/compat.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/compat.rs @@ -102,9 +102,7 @@ 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", - ); + 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 { @@ -1861,6 +1859,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 @@ -1873,7 +1881,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) {{", @@ -1892,7 +1901,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 {{", @@ -1903,6 +1913,7 @@ pub(super) fn render_compat_mbarrier_extended(catalog: &CatalogFile, hash: &str) } MbarrierExtendedOperation::FenceProxyAsyncSharedCta | MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster + | MbarrierExtendedOperation::FenceAcquireSharedClusterCluster | MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster | MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster => { writeln!(output, "pub unsafe fn {}() {{", record.rust.name).unwrap(); diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/warp_sync.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/warp_sync.rs index 217452d806..b8709ae8bb 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 @@ -787,7 +787,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 +801,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 +810,7 @@ fn verify_mbarrier_extended( ), MbarrierExtendedOperation::FenceProxyAsyncSharedCta | MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster + | MbarrierExtendedOperation::FenceAcquireSharedClusterCluster | MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster | MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster => ( 0, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/mma.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/mma.rs index ac5b1e5429..a6ac95879a 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 @@ -375,7 +375,7 @@ 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(), 1030); + assert_eq!(catalog.intrinsics.len(), 1033); let records: Vec<_> = register_mmas(&catalog).collect(); assert_eq!(records.len(), 155); let generated_records = records 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..5f2da1e1be 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/validate.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/validate.rs @@ -1105,6 +1105,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 +1143,10 @@ pub(super) fn validate_renderable(catalog: &CatalogFile) -> Result<()> { MbarrierExtendedAdapter::ZeroOperandsToVoid, ) => record.id == "fence_proxy_async_generic_acquire_shared_cluster_cluster", + ( + MbarrierExtendedOperation::FenceAcquireSharedClusterCluster, + MbarrierExtendedAdapter::ZeroOperandsToVoid, + ) => record.id == "fence_acquire_shared_cluster_cluster", ( MbarrierExtendedOperation::Nanosleep, MbarrierExtendedAdapter::NanosecondsToVoid, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mbarrier_extended.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mbarrier_extended.rs index 50800ed4e4..2c94b85121 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: "i1031", + 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: "i1032", + 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: "i1033", + id: "fence_acquire_shared_cluster_cluster", + operation_key: "fence.acquire.sync_restrict.shared_cluster.cluster", + source_record: Some("int_nvvm_fence_acquire_sync_restrict_space_cluster_scope_cluster"), + llvm_symbol: Some("llvm.nvvm.fence.acquire.sync_restrict.space.cluster.scope.cluster"), + ptx_native_instruction: None, + rust_arguments: &[], + rust_result: "()", + must_use: false, + dialect_op_type: "FenceAcquireSharedClusterClusterOp", + dialect_op_name: "nvvm.fence_acquire_shared_cluster_cluster", + dialect_operands: &[], + dialect_results: &[], + llvm_arguments: &[], + llvm_results: &[], + llvm_properties: &["IntrNoCallback"], + adapter: MbarrierExtendedAdapter::ZeroOperandsToVoid, + source_contract: MbarrierExtendedSourceContract::LlvmImported, + execution_scope: "cluster", + minimum_ptx: "8.6", + minimum_sm: "sm_90", + ptx_result: "()", + expected_ptx: instruction( + &[ + "fence", + "acquire", + "sync_restrict::shared::cluster", + "cluster", + ], + vec![], + ), + inline_ptx: "fence.acquire.sync_restrict::shared::cluster.cluster;", + inline_constraints: "~{memory}", + ptx_isa_section: "Parallel Synchronization and Communication Instructions: membar / fence", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar", + summary: "Acquires cluster-shared memory at cluster scope, without ordering global memory.", + }, MbarrierExtendedOperation::Nanosleep => MbarrierExtendedRecipe { operation, abi_id: "i0316", @@ -560,6 +697,9 @@ pub(in crate::resolve) fn expand_mbarrier_extended_admission( MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster, MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster, MbarrierExtendedOperation::Nanosleep, + MbarrierExtendedOperation::ArriveRemoteClusterRelaxed, + MbarrierExtendedOperation::TryWaitParityClusterRelaxed, + MbarrierExtendedOperation::FenceAcquireSharedClusterCluster, ]); let actual_operations = admission .variants @@ -658,7 +798,12 @@ pub(in crate::resolve) fn expand_mbarrier_extended_admission( admission.llvm_evidence_profile.clone() } IntrinsicBackend::LibNvvm => { - admission.libnvvm_evidence_profile.clone() + if matches!(recipe.operation, MbarrierExtendedOperation::ArriveRemoteClusterRelaxed | MbarrierExtendedOperation::TryWaitParityClusterRelaxed | MbarrierExtendedOperation::FenceAcquireSharedClusterCluster) { + admission.relaxed_libnvvm_evidence_profile.as_ref() + .unwrap_or(&admission.libnvvm_evidence_profile).clone() + } else { + admission.libnvvm_evidence_profile.clone() + } } }, targets: None, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mma/validate.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mma/validate.rs index d0adcd9a6c..f74d0e4c76 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mma/validate.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mma/validate.rs @@ -674,7 +674,10 @@ pub(in crate::resolve) fn validate_register_mma_block_scaled_policy( && policy.llvm_results == results && policy.ptx_result == "[f32; 4]" && mma.shape == shape - && matches!(mma.kind, Some(RegisterMmaKind::Mxf8f6f4 | RegisterMmaKind::Mxf4)) + && matches!( + mma.kind, + Some(RegisterMmaKind::Mxf8f6f4 | RegisterMmaKind::Mxf4) + ) && mma.operation == RegisterMmaOperation::Multiply && mma.accumulator == RegisterMmaAccumulator::F32 && mma.a_layout == RegisterMmaLayout::Row 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 bc68fbf623..750bbcf4a4 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/targets.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/targets.rs @@ -354,7 +354,11 @@ pub(super) fn validate_selected_target_predicates( "register_mma" => ( overlay_ptx == 87, policy.targets - == if policy.register_mma.as_ref().is_some_and(|mma| mma.kind == Some(RegisterMmaKind::Mxf4)) { + == 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" 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 c9bbbdc9ff..38ee20e217 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 @@ -1868,6 +1868,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 +1880,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 bbe5c2c982..ac2b75f197 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 @@ -251,16 +251,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, "i1031"); + assert_eq!(wait.abi_id, "i1032"); + 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, "i1033"); + 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([1031, 1032, 1033]) .map(|id| format!("i{id:04}")) .collect::>() ); 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 9dc54da8fc..34fbaabe53 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 @@ -49,7 +49,7 @@ fn overlay_manifest_loads_sorted_family_shards() { read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); assert_eq!(overlay.schema, OVERLAY_SCHEMA); assert_eq!(overlay.shards.len(), 67); - assert_eq!(overlay.intrinsics.len(), 1030); + assert_eq!(overlay.intrinsics.len(), 1033); assert_eq!( overlay .intrinsics 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 ad4ff3caf6..281e96ed9a 100644 --- a/cuda-oxide/crates/cuda-intrinsics/src/generated/abi_v1.rs +++ b/cuda-oxide/crates/cuda-intrinsics/src/generated/abi_v1.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Raw ABI functions recognized by cuda-oxide. @@ -3537,6 +3537,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 i1033() -> () { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033` 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 +5108,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 i1031(_arg0: u64) -> () { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1031` 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 +5270,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 i1032(_arg0: *const u64, _arg1: u32) -> bool { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1032` 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 , , ;`. diff --git a/cuda-oxide/crates/cuda-intrinsics/src/generated/mod.rs b/cuda-oxide/crates/cuda-intrinsics/src/generated/mod.rs index 96ab0e404b..66b4039dd2 100644 --- a/cuda-oxide/crates/cuda-intrinsics/src/generated/mod.rs +++ b/cuda-oxide/crates/cuda-intrinsics/src/generated/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. #[doc(hidden)] @@ -58,6 +58,9 @@ pub mod barrier { pub use crate::__cuda_oxide_intrinsic_abi_v1::i0885 as barrier_cta_arrive; pub use crate::__cuda_oxide_intrinsic_abi_v1::i0886 as barrier_cta_arrive_aligned; pub use crate::__cuda_oxide_intrinsic_abi_v1::i1017 as mbarrier_arrive_no_complete; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1031 as mbarrier_arrive_cluster_relaxed; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1032 as mbarrier_try_wait_parity_cluster_relaxed; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1033 as fence_acquire_shared_cluster_cluster; } /// Generated `bf16` intrinsic source API. @@ -2719,6 +2722,12 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_fence_acquire_shared_cluster_cluster_reexports_abi_i1033() { + let public: unsafe fn() -> () = super::barrier::fence_acquire_shared_cluster_cluster; + let canonical: unsafe fn() -> () = super::__cuda_oxide_intrinsic_abi_v1::i1033; + 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; @@ -3432,6 +3441,12 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_mbarrier_arrive_cluster_relaxed_reexports_abi_i1031() { + let public: unsafe fn(u64) -> () = super::barrier::mbarrier_arrive_cluster_relaxed; + let canonical: unsafe fn(u64) -> () = super::__cuda_oxide_intrinsic_abi_v1::i1031; + 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; @@ -3496,6 +3511,14 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_mbarrier_try_wait_parity_cluster_relaxed_reexports_abi_i1032() { + 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::i1032; + 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; diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/cutlass_backend.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/cutlass_backend.rs index 55a61c964f..a732593784 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/cutlass_backend.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/cutlass_backend.rs @@ -63,6 +63,8 @@ pub struct CutlassBackendConfig { 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 { @@ -71,10 +73,21 @@ impl CutlassBackendConfig { 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 { @@ -162,6 +175,7 @@ pub(crate) fn compile_to_cubin( )); } parse_sm_target(target)?; + let pipeline = precompiled_pipeline(config)?; let compiler = CutlassCompilerLibrary::load(&config.compiler_library, CUTLASS_COMPILER_LIBRARY_SHA256) @@ -174,6 +188,8 @@ pub(crate) fn compile_to_cubin( &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}")))?; @@ -182,7 +198,7 @@ pub(crate) fn compile_to_cubin( } let object = compiler - .compile_precompiled_mlir_to_object(&execution_mlir, target, CUTLASS_PRECOMPILED_PIPELINE) + .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}")))?; @@ -208,8 +224,10 @@ fn add_execution_manifest( compiler_sha256: &str, target: &str, profile: &str, + pipeline: &str, + optimization_level: u32, ) -> Result { - let pipeline_sha256 = format!("{:x}", Sha256::digest(CUTLASS_PRECOMPILED_PIPELINE)); + let pipeline_sha256 = format!("{:x}", Sha256::digest(pipeline)); let manifest = json!({ "profile": profile, "target": target, @@ -218,6 +236,7 @@ fn add_execution_manifest( "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!( @@ -1045,6 +1064,19 @@ mod tests { 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, 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 afc3b3fcf4..8fb9842169 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/active_mask.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/active_mask.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 b7caa8aa30..4f605e96d1 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/clc.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/clc.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 3430eb6925..ae5dbc1423 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cluster_barrier.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cluster_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 98760ed4a1..790cb3ada5 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cluster_memory.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cluster_memory.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 2a189f836f..4f36138269 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cp_async.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cp_async.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 dc3d0e2603..11637999a9 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/debug_control.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/debug_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 019b2055ba..4e94416794 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/dotprod.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/dotprod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 d374b82cea..834828032f 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/elect.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/elect.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 6c50be3482..edb33cdcd4 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/execution_control.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/execution_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 16d7e0c6ac..43f92c9e0c 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/extended_minmax.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/extended_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 f4c86981e0..304f8d01de 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/integer_minmax.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/integer_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 cb69157d3f..32a00971ac 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/ldmatrix.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/ldmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 f86b2989eb..190b08663e 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mbarrier_basic.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mbarrier_basic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 a3f448f20c..6a0421b318 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mbarrier_extended.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mbarrier_extended.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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:i1031", + id: "mbarrier_arrive_cluster_relaxed", + abi_id: "i1031", + 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:i1032", + id: "mbarrier_try_wait_parity_cluster_relaxed", + abi_id: "i1032", + 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:i1033", + id: "fence_acquire_shared_cluster_cluster", + abi_id: "i1033", + 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 7157a97038..e83c26a1c2 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mod.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated target requirements and separately imported LLVM/selection facts. 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 b0f8f2626b..8cbc73bf18 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/movmatrix.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 ca4cc1bae2..7338e9413e 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_alu.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_alu.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 e5459ff4d9..c009b74bd0 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_atomic.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_atomic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 85576dd3f2..65a54a8d29 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_conversion.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `packed_conversion` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/prmt.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/prmt.rs index f08177d609..af7c0c9a71 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/prmt.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 c80a5fb56d..21a3ee59a4 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/redux.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/redux.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 1095c95698..250ecc4d58 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/register_mma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/register_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `register_mma` intrinsics. 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 ff35fb2bdb..bbed01faff 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_arithmetic.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_arithmetic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 4da262f52b..b906823d10 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_conversion.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 e251582e2a..3a63bf0538 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_math.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_math.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 8b306e6f78..2c2b118e58 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sparse_mma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 0baaa4fff4..143918e7a4 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sreg.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 c6ba9edc19..2da548ffab 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/stmatrix.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 5ee60ae9d8..28753d72d4 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sync.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sync.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 0be529c4bd..368f56d613 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_cp.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_cp.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 b91bbb65d1..f93e5ce7b1 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_ld.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_ld.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 1296ccb9bd..23847aded5 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_mma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 7b31596e3b..6e5ddaf02e 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_mma_ws.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_mma_ws.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 ebbc67704b..39a0535a34 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_other.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_other.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 a42f54ce4e..babb9215b9 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_st.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_st.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 d380cc3279..b44d67c960 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/active_mask.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/active_mask.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 223e8a5050..e0efe95b3f 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/clc.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/clc.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 09238e1680..9a47f844fa 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cluster_barrier.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cluster_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 f2d96a889b..92f3ff9d23 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cluster_memory.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cluster_memory.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 2c0a3ff378..47323504ba 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cp_async.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cp_async.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 5c566d6261..cec3b83343 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/debug_control.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/debug_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 8d3ba3b38b..f4cb787fa9 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/dotprod.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/dotprod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 7ca2e2ee58..54d0376bd8 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/elect.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/elect.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 19fd2c48f3..f66431a9bf 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/execution_control.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/execution_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 43ec7aa177..8e5e21ea7a 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/extended_minmax.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/extended_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 dbc1520fb0..3b6dd0ea0a 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/integer_minmax.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/integer_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 3aaaf1c345..f9926ddba1 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/ldmatrix.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/ldmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 cef3e0fa35..151d4cd152 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mbarrier_basic.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mbarrier_basic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 c098e406f9..f95abad6b8 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mbarrier_extended.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mbarrier_extended.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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:i1031").unwrap(); + assert_eq!(target.id, "mbarrier_arrive_cluster_relaxed"); + assert_eq!(target.abi_id, "i1031"); + 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:i1032").unwrap(); + assert_eq!(target.id, "mbarrier_try_wait_parity_cluster_relaxed"); + assert_eq!(target.abi_id, "i1032"); + 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:i1033").unwrap(); + assert_eq!(target.id, "fence_acquire_shared_cluster_cluster"); + assert_eq!(target.abi_id, "i1033"); + 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 7d52a9a9cf..ee8687143b 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mod.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 f3f76a35a2..85d2f33617 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/movmatrix.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 0994743117..299909a600 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_alu.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_alu.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 4350ca8393..61eb9bafd9 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_atomic.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_atomic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 192ef263c9..b0eca3761b 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_conversion.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `packed_conversion` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/prmt.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/prmt.rs index 5b9463fa30..3bada32df1 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/prmt.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 b78c91a47f..e748b4dbdf 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/redux.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/redux.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 3e02995e4d..31ebcd6c2b 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/register_mma_float.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/register_mma_float.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `register_mma_float` intrinsics. 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 ee7c090e63..92e7f35986 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/register_mma_s32.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/register_mma_s32.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 eba216e845..c67ccc4deb 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_arithmetic.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_arithmetic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 3a73179ff6..d5f28dbfa9 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_conversion.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 5c129acd98..6014e50316 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_math.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_math.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 f1dad48bf5..45560c72b9 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sparse_mma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 3dfd084cd5..6e2ddd2a55 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sreg.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 fcbd86e25d..a1201489b1 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/stmatrix.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 61e475e20e..f06d8b1d0f 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sync.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sync.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 b99750561b..d282a24f1b 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_cp.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_cp.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 d34bfd8e1c..7a19022cc5 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_ld.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_ld.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 05c3ec1b68..f5bd251b50 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_mma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 543cc01da0..0361932fae 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_mma_ws.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_mma_ws.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 5aa5c9ef05..a7d8a83af5 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_other.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_other.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 2350c3f522..bc7c3aed5a 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_st.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_st.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 40a7787339..5fdffb1846 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `tma` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/vote.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/vote.rs index c2c543531c..d0c876f10c 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/vote.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/vote.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 e82ac6f7ae..161c815bad 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_barrier.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 fc17c2033d..dabd264cef 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_match.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_match.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 fed76e5b36..f1d362d05c 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_shuffle.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_shuffle.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 a37aa2d869..f0d005fc61 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/wgmma_control.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 0c8045c8c4..07f6fa3c0a 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `tma` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/vote.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/vote.rs index 1227cb99d2..bd7cc306d8 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/vote.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/vote.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 73bb06d039..c71dc347b0 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_barrier.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 a3314d11f5..dce4ea9e7f 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_match.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_match.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 56e3890f17..33a1a4ec3e 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_shuffle.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_shuffle.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 37408a8801..74df79779b 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/wgmma_control.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `wgmma_control` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/prep.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/prep.rs index 74a5c6388a..da7405da7a 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/prep.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/prep.rs @@ -246,6 +246,9 @@ fn has_high_level_cute_operations(ctx: &Context, root: Ptr) -> bool { || 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) @@ -288,6 +291,18 @@ fn has_high_level_cute_operations(ctx: &Context, root: Ptr) -> bool { || 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) @@ -779,6 +794,26 @@ mod tests { 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::(); 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 index 78e7bbdad2..b05f2805c1 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/tcgen05.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/tcgen05.rs @@ -21,13 +21,14 @@ use dialect_nvvm::ops::{ BarrierCtaSyncCountOp, ClusterBarrierModeAttr, ClusterBarrierOp, CpAsyncBulkCommitGroupOp, CpAsyncBulkTensorG2sTile2dMulticastCg2Op, CpAsyncBulkTensorS2gTile2dOp, CpAsyncBulkWaitGroupReadOp, CvtaGenericToSharedOffsetOp, ElectSyncOp, - FenceMbarrierInitReleaseClusterOp, FenceProxyAsyncSharedCtaOp, MapaSharedClusterOp, - MbarrierArriveClusterOp, MbarrierArriveExpectTxClusterOp, MbarrierInitSharedOp, - MbarrierTryWaitParityClusterOp, PrefetchTensorMapOp, Tcgen05AllocCg2Op, - Tcgen05CommitMulticastCg2Op, Tcgen05DeallocCg2Op, Tcgen05FenceAfterThreadSyncOp, - Tcgen05FenceBeforeThreadSyncOp, Tcgen05Ld32x32bX32RawOp, Tcgen05LoadWaitOp, - Tcgen05MmaCollectorAAttr, Tcgen05MmaCtaGroupAttr, Tcgen05MmaFormAttr, Tcgen05MmaKindAttr, - Tcgen05MmaOp, Tcgen05RelinquishAllocPermitCg2Op, + 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, @@ -164,6 +165,303 @@ fn cutlass_tcgen05_pack_compiles() { ); } +#[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 = pliron_mlir_export::render_module(&target); + 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. @@ -313,6 +611,24 @@ fn cutlass_cluster_tma_pack_compiles() { 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:i1031"), + ); + let relaxed_wait = append::( + &mut ctx, + block, + vec![a[2], zero], + vec![i1], + Some("v1:i1032"), + ); + 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:i1033")); append::(&mut ctx, block, vec![], vec![], None); let profile = CutlassFullCuteMlir22::new("sm_100a").unwrap(); @@ -320,11 +636,41 @@ fn cutlass_cluster_tma_pack_compiles() { let mlir = pliron_mlir_export::render_module(&target); 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", CUTLASS_PRECOMPILED_PIPELINE) + .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" @@ -345,3 +691,638 @@ fn integer_constant( ); 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 = pliron_mlir_export::render_module(&target); + 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 = pliron_mlir_export::render_module(&result); + 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 = pliron_mlir_export::render_module(&target); + 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 = pliron_mlir_export::render_module(&target); + 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/src/cute.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute.rs index 1871774afe..45f852faf5 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute.rs @@ -66,6 +66,10 @@ pub(crate) fn register_cute_elementwise_mappings( // 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)?; @@ -996,3 +1000,77 @@ fn u64_bits(value: u64) -> i128 { fn stringify(error: impl std::fmt::Display) -> String { error.to_string() } + +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_sm100.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_sm100.rs index a278d6ba4b..d90ba2eab1 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_sm100.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_sm100.rs @@ -436,19 +436,23 @@ fn emit_pipeline( Ok(()) } -struct Emit<'a, 'b> { - session: &'a mut TranslationSession<'b>, - location: MlirLocation, - ops: Vec, +pub(crate) struct Emit<'a, 'b> { + pub(crate) session: &'a mut TranslationSession<'b>, + pub(crate) location: MlirLocation, + pub(crate) ops: Vec, } impl Emit<'_, '_> { - fn operation(&self, name: &str, operands: Vec) -> Result { + 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) } - fn fresh(&mut self, ty: MlirType) -> (MlirResult, MlirValueUse) { + pub(crate) fn fresh(&mut self, ty: MlirType) -> (MlirResult, MlirValueUse) { let id = self.session.fresh_value(); (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) } @@ -457,7 +461,7 @@ impl Emit<'_, '_> { self.ops.push(op); Ok(()) } - fn value( + pub(crate) fn value( &mut self, name: &str, operands: Vec, @@ -469,7 +473,7 @@ impl Emit<'_, '_> { self.ops.push(op); Ok(v) } - fn cast( + pub(crate) fn cast( &mut self, name: &str, operand: MlirValueUse, @@ -477,7 +481,7 @@ impl Emit<'_, '_> { ) -> Result { self.value(name, vec![operand], ty) } - fn binary( + pub(crate) fn binary( &mut self, name: &str, a: MlirValueUse, @@ -486,7 +490,7 @@ impl Emit<'_, '_> { let ty = a.ty.clone(); self.value(name, vec![a, b], ty) } - fn constant(&mut self, bits: u32, value: i128) -> Result { + 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![])?; @@ -496,7 +500,11 @@ impl Emit<'_, '_> { self.ops.push(op); Ok(v) } - fn llvm_ptr(&mut self, ptr: MlirValueUse, space: u32) -> Result { + pub(crate) fn llvm_ptr( + &mut self, + ptr: MlirValueUse, + space: u32, + ) -> Result { let ty = llvm_ptr_type(space)?; if ptr.ty == ty { Ok(ptr) @@ -504,7 +512,7 @@ impl Emit<'_, '_> { self.cast("llvm.addrspacecast", ptr, ty) } } - fn cute_ptr( + pub(crate) fn cute_ptr( &mut self, ptr: MlirValueUse, element: &str, @@ -628,19 +636,19 @@ impl Emit<'_, '_> { fn dialect(s: &str) -> Result { MlirType::dialect(s) } -fn llvm_ptr_type(space: u32) -> Result { +pub(crate) fn llvm_ptr_type(space: u32) -> Result { dialect(&if space == 0 { "!llvm.ptr".into() } else { format!("!llvm.ptr<{space}>") }) } -fn property(op: &mut MlirOperation, key: &str, value: &str) -> Result<(), String> { +pub(crate) fn property(op: &mut MlirOperation, key: &str, value: &str) -> Result<(), String> { op.properties .insert(key.into(), MlirAttribute::dialect(value)?); Ok(()) } -fn integer_property(op: &mut MlirOperation, key: &str, bits: u32, value: i128) { +pub(crate) fn integer_property(op: &mut MlirOperation, key: &str, bits: u32, value: i128) { op.properties.insert( key.into(), MlirAttribute::Integer { @@ -649,7 +657,7 @@ fn integer_property(op: &mut MlirOperation, key: &str, bits: u32, value: i128) { }, ); } -fn segments(op: &mut MlirOperation, values: &[i32]) { +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_sm100_attention.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_sm100_attention.rs new file mode 100644 index 0000000000..cd1576cb47 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_sm100_attention.rs @@ -0,0 +1,270 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Independent layout-aware SM100 MMA and tiled TMA operations. +use crate::cute_sm100::{Emit, integer_property, llvm_ptr_type, property, segments}; +use dialect_cute::{attributes::CuteComposedLayoutAttr, sm100_attention_ops::*}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; +use pliron_mlir_export::{ + DropAttribute, MlirOperation, MlirType, MlirValueUse, OperationInput, OperationTranslation, + TranslationError, TranslationRegistry, TranslationSession, +}; + +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(()) +} +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(()) +} +struct Mma(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 + } +} +fn descriptor( + e: &mut Emit<'_, '_>, + ptr: MlirValueUse, + byte_offset: MlirValueUse, + layout: &CuteComposedLayoutAttr, + rows: u32, + k: u32, + mn: bool, +) -> Result { + let sw = layout.0.outer(); + let ptr = e.cute_ptr( + ptr, + "f16", + "smem", + 128, + Some(format!("S<{},{},{}>", sw.bits, sw.base + 1, sw.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) +} +struct Tma(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/lib.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/lib.rs index 609ddfd562..9eecaf9988 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/lib.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/lib.rs @@ -13,6 +13,7 @@ mod cute_attention; mod cute_gemm; mod cute_gemv; mod cute_sm100; +mod cute_sm100_attention; mod grid_constant; mod mir_checked; mod mir_core; @@ -20,8 +21,10 @@ mod mir_memory; mod mir_shared; mod nvvm_attention; mod nvvm_cluster; +mod nvvm_numeric; mod nvvm_sm120; mod nvvm_sregs; +mod nvvm_sync_restrict; mod nvvm_tcgen05; mod packs; mod primitives; @@ -35,8 +38,10 @@ pub(crate) use cute_sm100::register_cute_sm100_pack; pub use mir_core::register_mir_core_pack; pub use nvvm_attention::register_nvvm_attention_pack; pub use nvvm_cluster::register_nvvm_cluster_pack; +pub use nvvm_numeric::register_nvvm_numeric_pack; pub use nvvm_sm120::register_nvvm_sm120_pack; pub use nvvm_sregs::register_nvvm_sreg_pack; +pub use nvvm_sync_restrict::register_nvvm_sync_restrict_pack; pub use nvvm_tcgen05::register_nvvm_tcgen05_pack; pub use packs::register_builtin_pack; pub use primitives::{CutlassPrimitivesMlir22, verify_primitives_source}; 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 index 30923c99df..d3afd1a46d 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir_core.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir_core.rs @@ -144,6 +144,7 @@ impl OperationTranslation for FunctionTranslation { .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)?; } @@ -218,6 +219,28 @@ fn move_exact_block_contract(target: &mut MlirOperation) -> Result<(), String> { } } +/// 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"] 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 index 57010fdb8d..7ff9ddb4e6 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir_memory.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir_memory.rs @@ -708,23 +708,77 @@ impl OperationTranslation for CallTranslation { if callee.is_empty() || callee.chars().any(char::is_whitespace) { return Err(format!("mir.call has invalid callee {callee:?}")); } - if callee == dialect_mir::rust_intrinsics::CALLEE_MAXNUM_NSZ_F32 { - if input.operands.len() != 2 || input.results.len() != 1 { - return Err("Rust f32::max requires two operands and one result".into()); + // 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)?]); } - let f32_type = MlirType::Float(MlirFloatType::F32); - if input.operands.iter().any(|operand| operand.ty != f32_type) - || input.results[0].ty != f32_type - { - return Err("Rust f32::max operands and result must all be f32".into()); + 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]); } - // CUTLASS 4.7 cute/arch/nvvm_wrappers.py::fmax documents this exact - // modifier-free nvvm.fmax as maximumNumber, preserving subnormals. - // Unlike llvm.maxnum, PTX max.f32 ignores a quiet OR signaling NaN - // when the other operand is numeric. Its +0 preference is allowed - // by core::intrinsics::maximum_number_nsz_f32's signed-zero contract. - // Do not add nan (NaN propagation), ftz, abs, or fast-math flags. - return Ok(vec![renamed("nvvm.fmax", input)?]); + 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() 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 index 90d4746f1d..fa174df3a5 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_attention.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_attention.rs @@ -30,6 +30,12 @@ pub fn register_nvvm_attention_pack( registry.register_attribute::(DropAttribute)?; registry.register_operation::(MathTranslation)?; registry.register_operation::(ShuffleTranslation)?; + register_nvvm_async_control_pack(registry) +} + +pub(crate) fn register_nvvm_async_control_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { registry.register_operation::(CommitTranslation)?; registry.register_operation::(WaitTranslation)?; Ok(()) @@ -42,10 +48,17 @@ impl OperationTranslation for MathTranslation { ctx: &Context, source: Ptr, mut input: OperationInput, - _session: &mut TranslationSession<'_>, + session: &mut TranslationSession<'_>, ) -> Result, String> { let math = ScalarMathOp::new(source); math.verify(ctx).map_err(|e| e.to_string())?; + // Log2 uses the common numeric mapping; exp2 and reciprocal retain + // the SM120 attention schedule's validated CUTLASS lowering below. + if math.get_attr_nvvm_scalar_math_operation(ctx).as_deref() + == Some(&ScalarMathOperationAttr::Lg2) + { + 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), 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 index 823fc1c5bf..2e45f7e71d 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_cluster.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_cluster.rs @@ -120,6 +120,11 @@ pub fn register_nvvm_cluster_pack( NoOperands("nvvm.fence.mbarrier.init") ); register!(FenceProxyAsyncSharedCtaOp, "v1:i0312", AsyncFence); + register!( + FenceAcquireSharedClusterClusterOp, + "v1:i1033", + SharedAcquireFence + ); register!( BarrierCtaSyncCountOp, "v1:i0883", @@ -168,14 +173,37 @@ pub fn register_nvvm_cluster_pack( register!( MbarrierTryWaitParitySharedOp, "v1:i0310", - WaitParity { cluster: false } + WaitParity { + cluster: false, + relaxed: false + } ); register!( MbarrierTryWaitParityClusterOp, "v1:i0311", - WaitParity { cluster: true } + WaitParity { + cluster: true, + relaxed: false + } + ); + register!( + MbarrierArriveClusterOp, + "v1:i0308", + RemoteArrive { relaxed: false } + ); + register!( + MbarrierArriveClusterRelaxedOp, + "v1:i1031", + RemoteArrive { relaxed: true } + ); + register!( + MbarrierTryWaitParityClusterRelaxedOp, + "v1:i1032", + WaitParity { + cluster: true, + relaxed: true + } ); - register!(MbarrierArriveClusterOp, "v1:i0308", RemoteArrive); register!(MapaSharedClusterOp, "v1:i0320", MapShared); register!(ElectSyncOp, "v1:i0367", Elect); register!(PrefetchTensorMapOp, "v1:i0887", Prefetch); @@ -207,13 +235,14 @@ enum Kind { Sreg(&'static str), NoOperands(&'static str), AsyncFence, + SharedAcquireFence, CountedBarrier { arrive: bool, aligned: bool }, BarrierInit, BarrierInval, BarrierArrive, ExpectTx { cluster: bool }, - WaitParity { cluster: bool }, - RemoteArrive, + WaitParity { cluster: bool, relaxed: bool }, + RemoteArrive { relaxed: bool }, MapShared, Elect, Prefetch, @@ -250,6 +279,12 @@ impl OperationTranslation for Recipe { 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)?; @@ -303,7 +338,7 @@ impl OperationTranslation for Recipe { } op } - Kind::WaitParity { cluster } => { + Kind::WaitParity { cluster, relaxed } => { signature(&input, &[Pointer(&[0, 3]), Int(32)], &[Int(1)])?; let mut op = operation( "nvvm.mbarrier.wait.parity", @@ -314,11 +349,19 @@ impl OperationTranslation for Recipe { dialect_property(&mut op, "kind", "#nvvm.mbar_wait")?; if cluster { dialect_property(&mut op, "scope", "#nvvm.mbar_scope")?; - dialect_property(&mut op, "order", "#nvvm.mem_order")?; + dialect_property( + &mut op, + "order", + if relaxed { + "#nvvm.mem_order" + } else { + "#nvvm.mem_order" + }, + )?; } op } - Kind::RemoteArrive => { + Kind::RemoteArrive { relaxed } => { signature(&input, &[Int(64)], &[])?; let pointer = cast( session, @@ -330,6 +373,10 @@ impl OperationTranslation for Recipe { )?; 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 => { @@ -899,6 +946,24 @@ mod tests { } } + #[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()); + } + } + #[test] fn barrier_invalidation_preserves_pointer_and_has_no_hidden_sync() { for argument in [0, 1] { @@ -1019,6 +1084,40 @@ mod tests { assert!(!text.contains("\"nvvm.mbarrier.try_wait.parity\"")); } + #[test] + fn relaxed_cluster_arrival_and_wait_keep_scope_without_acquire_release() { + let mut fixture = Fixture::new(); + fixture.append::("v1:i1031", vec![fixture.arg(6)], vec![]); + fixture.append::( + "v1:i1032", + 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}"); + } + + #[test] + fn restricted_acquire_fence_retains_shared_only_ordering() { + let mut fixture = Fixture::new(); + fixture.append::("v1:i1033", 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}"); + } + #[test] fn elect_preserves_both_lane_and_predicate_results() { let mut fixture = Fixture::new(); 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..8d425d0fdf --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_numeric.rs @@ -0,0 +1,515 @@ +/* + * 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 dialect_nvvm::ops::*; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; +use pliron_mlir_export::{ + DropAttribute, MlirAttribute, MlirFloatType, MlirOperation, MlirType, OperationInput, + OperationTranslation, TranslationError, TranslationRegistry, TranslationSession, +}; + +type VerifyOperation = fn(&Context, Ptr) -> Result<(), String>; + +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) + } +} + +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)?; + 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); + + 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(()) +} + +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) +} + +struct IntrinsicTranslation { + marker: &'static str, + name: &'static str, + operands: usize, + shuffle: bool, + 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) + } +} + +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 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") + } + _ => { + 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)?]) + } +} + +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 { + marker: &'static str, + name: &'static str, + operands: usize, + 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) + } +} + +struct RegisterTranslation(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]) + } +} +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(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_sm100::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_sm100::property(&mut op, "kind", "#nvvm")?; + } + 1 => { + crate::cute_sm100::integer_property(&mut op, "size", 32, 4); + crate::cute_sm100::property(&mut op, "modifier", "#nvvm")?; + } + 2 => { + op.properties + .insert("noinc".into(), MlirAttribute::Bool(true)); + } + 3 => { + crate::cute_sm100::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 index 004a2e86d6..d0fdb59947 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_sm120.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_sm120.rs @@ -32,7 +32,6 @@ pub fn register_nvvm_sm120_pack( registry.register_attribute::(DropAttribute)?; registry.register_operation::(LoadTranslation)?; registry.register_operation::(StoreTranslation)?; - registry.register_operation::(PackTranslation)?; registry.register_operation::(TmaTranslation)?; registry.register_attribute::(DropAttribute)?; registry.register_attribute::(DropAttribute)?; @@ -271,53 +270,6 @@ impl OperationTranslation for StoreTranslation { } } -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)?, - ]) - } -} - struct TmaTranslation; impl OperationTranslation for TmaTranslation { fn translate( 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..8e43d67406 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_sync_restrict.rs @@ -0,0 +1,102 @@ +// 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 dialect_nvvm::ops::{ + FenceProxyAsyncGenericAcquireSharedClusterClusterOp, + FenceProxyAsyncGenericReleaseSharedCtaClusterOp, +}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; +use pliron_mlir_export::{ + MlirAttribute, MlirOperation, OperationInput, OperationTranslation, TranslationError, + TranslationRegistry, TranslationSession, +}; + +/// Register cluster-scoped acquire/release ordering between generic and async proxies. +pub fn register_nvvm_sync_restrict_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + 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(()) +} + +struct ProxyFenceTranslation { + marker: &'static str, + order: &'static str, + 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 index 1370a7ca73..7db962bb7c 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_tcgen05.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_tcgen05.rs @@ -10,13 +10,8 @@ //! 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::{ - CvtaGenericToClusterSharedOffsetOp, CvtaGenericToSharedOffsetOp, Tcgen05AllocCg2Op, - Tcgen05CommitMulticastCg2Op, Tcgen05DeallocCg2Op, Tcgen05FenceAfterThreadSyncOp, - Tcgen05FenceBeforeThreadSyncOp, Tcgen05Ld32x32bX32RawOp, Tcgen05LoadWaitOp, - Tcgen05MmaBBufferAttr, Tcgen05MmaBUsageAttr, Tcgen05MmaCollectorAAttr, Tcgen05MmaCtaGroupAttr, - Tcgen05MmaFormAttr, Tcgen05MmaKindAttr, Tcgen05MmaOp, Tcgen05RelinquishAllocPermitCg2Op, -}; +use dialect_nvvm::ops::*; + use pliron::{ common_traits::Verify, context::{Context, Ptr}, @@ -48,22 +43,118 @@ pub fn register_nvvm_tcgen05_pack( registry.register_operation::(SharedOffsetTranslation(false))?; registry .register_operation::(SharedOffsetTranslation(true))?; - registry.register_operation::(LoadTranslation)?; - registry.register_operation::(PrimitiveTranslation(Primitive::Alloc))?; - registry.register_operation::(PrimitiveTranslation(Primitive::Dealloc))?; - registry.register_operation::(PrimitiveTranslation( + 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, - ))?; - registry.register_operation::(PrimitiveTranslation( + 2, + 2 + ); + primitive!( + Tcgen05RelinquishAllocPermitOp, + "v1:i0345", Primitive::Relinquish, - ))?; - registry.register_operation::(PrimitiveTranslation(Primitive::LoadWait))?; - registry.register_operation::(PrimitiveTranslation( + 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, - ))?; - registry.register_operation::(PrimitiveTranslation( + 1, + 0 + ); + primitive!( + Tcgen05FenceAfterThreadSyncOp, + "v1:i0347", Primitive::FenceAfter, - ))?; + 1, + 0 + ); + 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(()) } @@ -133,15 +224,26 @@ impl OperationTranslation for MmaTranslation { ) -> Result, String> { let mma = Tcgen05MmaOp::new(source); mma.verify(ctx).map_err(|error| error.to_string())?; - if mma.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() != Some(&Tcgen05MmaFormAttr::Shared) - || mma.get_attr_nvvm_tcgen05_mma_kind(ctx).as_deref() != Some(&Tcgen05MmaKindAttr::F16) - || mma.get_attr_nvvm_tcgen05_mma_cta_group(ctx).as_deref() - != Some(&Tcgen05MmaCtaGroupAttr::Cg2) - { + 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 shared-A FP16 CTA-group-2 MMA only".into(), + "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", @@ -149,15 +251,29 @@ impl OperationTranslation for MmaTranslation { Some(Tcgen05MmaCollectorAAttr::Use) => "use", None => return Err("TCGen05 MMA requires a collector-A selector".into()), }; - require_input(&mut input, "v1:i0763", 5, 0)?; + 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()?); + target + .properties + .insert("ctaGroup".into(), cta_group_value(group)?); target.properties.insert( "collectorOp".into(), MlirAttribute::dialect(format!("#nvvm.tcgen05_mma_collectorop<{collector}>"))?, @@ -166,7 +282,8 @@ impl OperationTranslation for MmaTranslation { "operandSegmentSizes".into(), MlirAttribute::DenseI32Array(vec![1, 1, 1, 1, 1, 0, 0]), ); - Ok(vec![cast, target]) + operations.push(target); + Ok(operations) } } @@ -177,11 +294,20 @@ enum Primitive { Commit, Relinquish, LoadWait, + StoreWait, FenceBefore, FenceAfter, } -struct PrimitiveTranslation(Primitive); +type VerifyOperation = fn(&Context, Ptr) -> Result<(), String>; + +struct PrimitiveTranslation { + marker: &'static str, + kind: Primitive, + group: u32, + operands: usize, + verify: VerifyOperation, +} impl OperationTranslation for PrimitiveTranslation { fn translate( @@ -191,61 +317,16 @@ impl OperationTranslation for PrimitiveTranslation { mut input: OperationInput, session: &mut TranslationSession<'_>, ) -> Result, String> { - let (name, marker, operands) = match self.0 { - Primitive::Alloc => { - Tcgen05AllocCg2Op::new(source) - .verify(ctx) - .map_err(|error| error.to_string())?; - ("nvvm.tcgen05.alloc", "v1:i0359", 2) - } - Primitive::Dealloc => { - Tcgen05DeallocCg2Op::new(source) - .verify(ctx) - .map_err(|error| error.to_string())?; - ("nvvm.tcgen05.dealloc", "v1:i0360", 2) - } - Primitive::Commit => { - Tcgen05CommitMulticastCg2Op::new(source) - .verify(ctx) - .map_err(|error| error.to_string())?; - ("nvvm.tcgen05.commit", "v1:i0365", 2) - } - Primitive::Relinquish => { - Tcgen05RelinquishAllocPermitCg2Op::new(source) - .verify(ctx) - .map_err(|error| error.to_string())?; - ("nvvm.tcgen05.relinquish_alloc_permit", "v1:i0361", 0) - } - Primitive::LoadWait => { - Tcgen05LoadWaitOp::new(source) - .verify(ctx) - .map_err(|error| error.to_string())?; - ("nvvm.tcgen05.wait", "v1:i0357", 0) - } - Primitive::FenceBefore => { - Tcgen05FenceBeforeThreadSyncOp::new(source) - .verify(ctx) - .map_err(|error| error.to_string())?; - ("nvvm.tcgen05.fence", "v1:i0346", 0) - } - Primitive::FenceAfter => { - Tcgen05FenceAfterThreadSyncOp::new(source) - .verify(ctx) - .map_err(|error| error.to_string())?; - ("nvvm.tcgen05.fence", "v1:i0347", 0) - } - }; - require_input(&mut input, marker, operands, 0)?; + (self.verify)(ctx, source)?; + require_input(&mut input, self.marker, self.operands, 0)?; let mut operations = vec![]; - if matches!(self.0, Primitive::Dealloc) { + 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; } - // Native SMEM pointer arguments can be generic after MIR casts. - // Preserve the address while making the required shared space explicit. - if matches!(self.0, Primitive::Alloc | Primitive::Commit) { + if matches!(self.kind, Primitive::Alloc | Primitive::Commit) { input.operands[0] = shared_pointer( session, &mut operations, @@ -253,28 +334,42 @@ impl OperationTranslation for PrimitiveTranslation { &input.location, )?; } - let mut target = operation(name, input.operands, vec![], &input.location)?; - let kind = match self.0 { - Primitive::LoadWait => Some("#nvvm.tcgen05_wait"), - Primitive::FenceBefore => Some("#nvvm.tcgen05_fence"), - Primitive::FenceAfter => Some("#nvvm.tcgen05_fence"), - _ => None, + 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()?); + target + .properties + .insert("group".into(), cta_group_value(self.group)?); } operations.push(target); Ok(operations) } } -struct LoadTranslation; +/// Scalar carriers retain the native register order. The native instructions +/// are asynchronous; only explicit wait operations introduce completion waits. +struct MemoryTranslation { + marker: &'static str, + width: usize, + store: bool, + pack: bool, + verify: VerifyOperation, +} -impl OperationTranslation for LoadTranslation { +impl OperationTranslation for MemoryTranslation { fn translate( &self, ctx: &Context, @@ -282,50 +377,117 @@ impl OperationTranslation for LoadTranslation { mut input: OperationInput, session: &mut TranslationSession<'_>, ) -> Result, String> { - Tcgen05Ld32x32bX32RawOp::new(source) - .verify(ctx) - .map_err(|error| error.to_string())?; - require_input(&mut input, "v1:i0664", 1, 32)?; + (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]; - let (vector_result, vector) = fresh( - session, - MlirType::Vector { - shape: vec![32], - element: Box::new(MlirType::Integer(32)), - }, - ); - let mut load = operation( - "nvvm.tcgen05.ld", - vec![tmem], - vec![vector_result], - &input.location, - )?; - load.properties.insert( + // 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")?, ); - operations.push(load); - // `raw` is asynchronous: retain the caller's explicit tcgen05.wait. - // Do not insert a wait here and change the source scheduling policy. - 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]), + if self.pack { + target.properties.insert( + if self.store { "unpack" } else { "pack" }.into(), + MlirAttribute::Unit, ); - operations.push(extract); + } + if self.store { + operations.push(target); + } else { + operations.insert(1, target); } Ok(operations) } } -fn require_input( +pub(crate) fn require_input( input: &mut OperationInput, marker: &str, operands: usize, @@ -351,8 +513,8 @@ fn require_input( Ok(()) } -fn cta_group() -> Result { - MlirAttribute::dialect("#nvvm.cta_group") +fn cta_group_value(group: u32) -> Result { + MlirAttribute::dialect(format!("#nvvm.cta_group")) } fn tmem_pointer( @@ -393,12 +555,15 @@ fn shared_pointer( Ok(converted) } -fn fresh(session: &mut TranslationSession<'_>, ty: MlirType) -> (MlirResult, MlirValueUse) { +pub(crate) fn fresh( + session: &mut TranslationSession<'_>, + ty: MlirType, +) -> (MlirResult, MlirValueUse) { let id = session.fresh_value(); (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) } -fn operation( +pub(crate) fn operation( name: &str, operands: Vec, results: Vec, @@ -588,16 +753,305 @@ mod tests { mma.set_attr_nvvm_tcgen05_mma_kind(&ctx, Tcgen05MmaKindAttr::Tf32); let error = translate(&ctx, &module).unwrap_err(); assert!( - error.contains("shared-A FP16 CTA-group-2 MMA only"), + error.contains("dense shared-A or tensor-A FP16 MMA only"), "{error}" ); - mma.set_attr_nvvm_tcgen05_mma_kind(&ctx, Tcgen05MmaKindAttr::F16); - mma.set_attr_nvvm_tcgen05_mma_cta_group(&ctx, Tcgen05MmaCtaGroupAttr::Cg1); - let error = translate(&ctx, &module).unwrap_err(); - assert!( - error.contains("shared-A FP16 CTA-group-2 MMA only"), - "{error}" + } + + #[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}" + ); + } + } + } + + #[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(); + 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}"); + } } #[test] diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/packs.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/packs.rs index 765dc06b5b..30fbcba9c6 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/packs.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/packs.rs @@ -11,7 +11,7 @@ use pliron::{ BoolAttr, FPDoubleAttr, FPHalfAttr, FPSingleAttr, IdentifierAttr, IntegerAttr, OperandSegmentSizesAttr, StringAttr, TypeAttr, UnitAttr, }, - ops::ModuleOp, + ops::{ConstantOp, ModuleOp}, type_interfaces::FunctionTypeInterface, types::{FP16Type, FP32Type, FP64Type, FunctionType, IntegerType, UnitType}, }, @@ -29,6 +29,7 @@ use pliron_mlir_export::{ /// 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)))?; @@ -80,6 +81,51 @@ impl OperationTranslation for BuiltinModuleTranslation { } } +// Scalar constants created by canonicalization and loop-unrolling passes use +// builtin.constant rather than the source-facing mir.constant operation. +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]) + } +} + struct IntegerTypeTranslation; impl TypeTranslation for IntegerTypeTranslation { diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/primitives.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/primitives.rs index 87cde7479b..90b1352d5e 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/primitives.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/primitives.rs @@ -68,8 +68,9 @@ impl MlirConsumerProfile for CutlassPrimitivesMlir22 { 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)?; - crate::register_nvvm_attention_pack(&mut r)?; r.seal(); Ok(r) } diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/profile.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/profile.rs index f9140667d4..7442e906da 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/profile.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/profile.rs @@ -132,11 +132,13 @@ impl MlirConsumerProfile for CutlassFullCuteMlir22 { register_nvvm_tcgen05_pack(&mut registry)?; register_nvvm_cluster_pack(&mut registry)?; crate::register_nvvm_sm120_pack(&mut registry)?; - crate::register_nvvm_attention_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_sm100_attention::register_cute_sm100_attention_pack(&mut registry)?; registry.seal(); Ok(registry) } 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..bc04f409df --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_cute.rs @@ -0,0 +1,175 @@ +/* + * 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_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] +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 = pliron_mlir_export::render_module(&target); + 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..28fd4088ef --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_descriptor.rs @@ -0,0 +1,218 @@ +/* + * 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 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 pliron_mlir_export::{MlirAttribute, MlirBlock, MlirOperation, MlirType, MlirValueId}; +use std::collections::HashMap; + +struct Probe { + definitions: HashMap, + offset: MlirValueId, + descriptors: Vec, +} + +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 { + 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(), + } + } + + 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, + } + } +} + +#[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); + } + } + } +} + +#[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..fc11c3e4a0 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_primitives.rs @@ -0,0 +1,220 @@ +/* + * 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}, +}; + +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] +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 = pliron_mlir_export::render_module(&result); + 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 index 6fd99efd8e..d2098c79df 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/checked_storage.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/checked_storage.rs @@ -332,3 +332,108 @@ fn explicitly_returned_zero_sized_aggregates_keep_their_call_result() { assert!(!text.contains("\"ub.poison\""), "{text}"); } } + +#[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}"); + } +} + +#[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}"); +} + +#[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}" + ); +} + +#[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 index 416dc9b42e..b05ae2cebb 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs @@ -295,6 +295,8 @@ fn production_kernel_gets_one_device_module_and_the_six_argument_launch_abi() { StringAttr::new("true".into()), ); for (name, value) in [ + ("maxntid", 256), + ("minctasm", 1), ("reqntid_x", 256), ("reqntid_y", 1), ("reqntid_z", 1), @@ -388,6 +390,16 @@ fn production_kernel_gets_one_device_module_and_the_six_argument_launch_abi() { 256, 1, 1 ])) ); + assert_eq!( + kernel.attributes.get("nvvm.minctasm"), + Some(&pliron_mlir_export::MlirAttribute::Integer { + value: 1, + ty: pliron_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 pliron_mlir_export::MlirAttribute::Type(pliron_mlir_export::MlirType::Function { inputs, .. @@ -448,4 +460,19 @@ fn production_kernel_gets_one_device_module_and_the_six_argument_launch_abi() { "{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); + 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/nvvm_attention.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/nvvm_attention.rs index a3dd8ecc47..07eefd33b0 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/nvvm_attention.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/nvvm_attention.rs @@ -338,3 +338,82 @@ fn cp_async_wait_rejects_runtime_or_out_of_range_counts() { assert!(f.translate().is_err()); } } + +#[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/sync_restrict.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/sync_restrict.rs new file mode 100644 index 0000000000..7b1867b257 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/sync_restrict.rs @@ -0,0 +1,180 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +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, +}; +use pliron_mlir_export::render_module; + +struct Fixture { + ctx: Context, + module: ModuleOp, + release: Ptr, + acquire: Ptr, +} + +impl Fixture { + 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, + } + } + + 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.map(|m| render_module(&m)).map_err(|e| e.to_string()) + } +} + +#[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}"); + } +} + +#[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}" + ); + } +} + +#[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}"); +} + +#[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/dialect-nvvm/src/ops/generated/active_mask.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/active_mask.rs index c81f2db9d8..078385d829 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/active_mask.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/active_mask.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operation for the generated active warp mask. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/clc.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/clc.rs index 4013cf93a5..930c63c4b5 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/clc.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/clc.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated Cluster Launch Control operations. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_barrier.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_barrier.rs index 79803262e4..2e257a6ad0 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_barrier.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated operations for the closed cluster-barrier family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_memory.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_memory.rs index 69fd5439b1..94331d1270 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_memory.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_memory.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for cluster address mapping and remote shared reads. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cp_async.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cp_async.rs index 06b071bb98..875746b067 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cp_async.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cp_async.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for classic `cp.async` instructions. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/debug_control.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/debug_control.rs index 784f32d80a..2489ca1da3 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/debug_control.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/debug_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for generated PTX debug controls. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/dotprod.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/dotprod.rs index f62064f0e3..f5eb651103 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/dotprod.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/dotprod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for generated packed integer dot products. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/elect.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/elect.rs index 396ff5745d..91477f32de 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/elect.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/elect.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated warp leader-election operation. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/execution_control.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/execution_control.rs index d0abd0d57a..988ecf3363 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/execution_control.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/execution_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated counted-barrier, grid-dependency, and register-control operations. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/extended_minmax.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/extended_minmax.rs index bd0ad945a1..fe77c3f8d0 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/extended_minmax.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/extended_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for extended floating-point min/max. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/integer_minmax.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/integer_minmax.rs index d394d4ce55..86ef3ca678 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/integer_minmax.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/integer_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for generated extended integer min/max. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/ldmatrix.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/ldmatrix.rs index e91892bab4..1858f36781 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/ldmatrix.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/ldmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operation and compatibility carriers for generated `ldmatrix`. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_basic.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_basic.rs index 06bb36b4e9..6225baa78e 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_basic.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_basic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for the basic shared-memory mbarrier lifecycle. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_extended.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_extended.rs index ef1d7c105c..5b4c8802b9 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_extended.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_extended.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for extended mbarrier and async-proxy instructions. @@ -120,6 +120,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 +301,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 +522,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 +596,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 34974e45fc..2361efb618 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mod.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. mod active_mask; diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/movmatrix.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/movmatrix.rs index 80f036089f..dfed79b641 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/movmatrix.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operation for generated in-register matrix transpose. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_alu.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_alu.rs index bf5292e9e2..22fdcbdf00 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_alu.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_alu.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for generated packed floating-point arithmetic. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_atomic.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_atomic.rs index 6a4332c4d3..d1b9ce2785 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_atomic.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_atomic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for the closed generated packed-atomic family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_conversion.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_conversion.rs index f74910caba..0819b068b1 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_conversion.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operation for generated packed conversion. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/prmt.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/prmt.rs index f0f19b48af..050e11f980 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/prmt.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for the closed generated `prmt` family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/redux.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/redux.rs index 94fc1e7a78..02c9eea5f1 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/redux.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/redux.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural NVVM operations for the closed generated `redux.sync` family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/register_mma.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/register_mma.rs index 272a9aff2e..6d0fb3e9c6 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/register_mma.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/register_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for generated register-only `mma.sync` variants. 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 c54022a8ed..e277f00bdc 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_arithmetic.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_arithmetic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for generated scalar arithmetic. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_conversion.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_conversion.rs index e786f8b9e5..5e8c37e0ae 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_conversion.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for generated `f32` to TF32 conversions. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_math.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_math.rs index 79114f32e2..36bc8288ed 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_math.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_math.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for generated scalar math. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sparse_mma.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sparse_mma.rs index 4f104f7921..d5fa484dba 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sparse_mma.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for generated sparse MMA variants. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sreg.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sreg.rs index dffa9763fe..0f353dafcc 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sreg.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural NVVM operations for generated special-register reads. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/stmatrix.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/stmatrix.rs index fe174e6716..8878e538bb 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/stmatrix.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for the four generated `stmatrix` stores. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sync.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sync.rs index c1128929b1..f2bde59a38 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sync.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sync.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 45079d7ea5..2d3357c7d1 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tcgen05.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tcgen05.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated Tensor Core Generation 5 operations. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tma.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tma.rs index 3ef8cc9cc3..d71f71bd0f 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tma.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated Tensor Memory Accelerator operations. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/vote.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/vote.rs index f8faefa53a..2b2af793dd 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/vote.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/vote.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for the generated `vote.sync` family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_barrier.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_barrier.rs index 6affee381b..b279a1e64e 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_barrier.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operation for generated warp synchronization. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_match.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_match.rs index dcb1457ded..37b12998f1 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_match.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_match.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for the generated `match.sync` family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_shuffle.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_shuffle.rs index e423ff3fbe..e11410ae6c 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_shuffle.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_shuffle.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for the generated `shfl.sync` family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/wgmma_control.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/wgmma_control.rs index b1379f7d7d..edf8ea6125 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/wgmma_control.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for generated WGMMA controls. diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/emit.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/emit.rs index 2602c2eec7..297c905266 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/emit.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/emit.rs @@ -595,6 +595,49 @@ mod tests { (undef.deref(ctx).get_result(0), undef) } + #[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()); + } + #[test] fn tma_boundary_recovers_only_generic_or_existing_cta_shared_pointers() { let mut ctx = Context::new(); @@ -1321,3 +1364,102 @@ pub(crate) fn emit_copy_tma_s2g( }; 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)) +} + +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/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/mod.rs index 0cbaa81a63..cd3fad21ca 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/mod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/mod.rs @@ -34,6 +34,7 @@ pub(crate) mod pipeline; pub(crate) mod pipeline_emit; pub(crate) mod scheduler; pub(crate) mod scheduler_emit; +pub(crate) mod sm100_attention_emit; pub(crate) mod sm100_emit; pub(crate) mod smem_mma; pub(crate) mod smem_mma_emit; @@ -57,6 +58,7 @@ use rustc_public::ty::FnDef; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum CuteFn { Sm100(sm100_emit::Sm100Fn), + Attention(sm100_attention_emit::AttentionFn), MakeTensorRead, MakeTensorWrite, ZippedDivideRead, @@ -108,6 +110,7 @@ pub(crate) enum CuteFn { TmaStoreTail, LoadTile, StoreTile, + StoreGlobalNoAllocate, AssumeDiv, CopyG2S, LoadMatrixA, @@ -137,6 +140,39 @@ pub(crate) struct CuteFnSpec { /// One entry per recognized function, including its generic argument schema. pub(crate) const CUTE_FNS: &[CuteFnSpec] = &[ + CuteFnSpec { + canonical: "cute_rs::sm100_attention::__compiler::mma_ss", + func: CuteFn::Attention(sm100_attention_emit::AttentionFn::MmaSs), + schema: &[ + SubstKind::Type, + SubstKind::Type, + SubstKind::Const, + SubstKind::Const, + SubstKind::Const, + SubstKind::Const, + ], + }, + CuteFnSpec { + canonical: "cute_rs::sm100_attention::__compiler::mma_ts", + func: CuteFn::Attention(sm100_attention_emit::AttentionFn::MmaTs), + schema: &[ + SubstKind::Type, + SubstKind::Const, + SubstKind::Const, + SubstKind::Const, + SubstKind::Const, + ], + }, + CuteFnSpec { + canonical: "cute_rs::sm100_attention::__compiler::tma_load", + func: CuteFn::Attention(sm100_attention_emit::AttentionFn::TmaLoad), + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100_attention::__compiler::tma_store", + func: CuteFn::Attention(sm100_attention_emit::AttentionFn::TmaStore), + schema: &[SubstKind::Type, SubstKind::Const], + }, CuteFnSpec { canonical: "cute_rs::sm100::__compiler::sm100_tma_store", func: CuteFn::Sm100(sm100_emit::Sm100Fn::TmaStore), @@ -499,6 +535,11 @@ pub(crate) const CUTE_FNS: &[CuteFnSpec] = &[ func: CuteFn::StoreTile, schema: &[SubstKind::Type, SubstKind::Const], }, + CuteFnSpec { + canonical: "cute_rs::tile::store_global_no_allocate", + func: CuteFn::StoreGlobalNoAllocate, + schema: &[SubstKind::Type, SubstKind::Const], + }, CuteFnSpec { canonical: "cute_rs::tile::assume_div", func: CuteFn::AssumeDiv, @@ -573,6 +614,7 @@ enum CuteIdentity { /// body-translated instead of dispatched through [`CUTE_FNS`]. New modules /// therefore require an explicit audit and an update to the exact-list test. const CUTE_SOURCE_ABSTRACTION_MODULES: &[&str] = &[ + "arch", "block_scaled", "block_scaled_mma", "epilogue", @@ -580,6 +622,7 @@ const CUTE_SOURCE_ABSTRACTION_MODULES: &[&str] = &[ "pipeline", "scheduler", "sm100", + "sm100_attention", "tensor", "tiled_copy", ]; @@ -649,11 +692,29 @@ mod canonical_path_tests { assert_eq!(path, "cute_rs::cute_rs::tile::load_tile"); } + #[test] + fn cache_aware_global_store_has_an_exact_compiler_boundary() { + let spec = super::CUTE_FNS + .iter() + .find(|spec| spec.canonical == "cute_rs::tile::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] + ); + } + #[test] fn source_abstraction_modules_are_exact_and_closed() { + assert!(is_source_abstraction_path( + "cute_rs::arch::assume_warp_uniform", + "cute_rs" + )); assert_eq!( super::CUTE_SOURCE_ABSTRACTION_MODULES, &[ + "arch", "block_scaled", "block_scaled_mma", "epilogue", @@ -661,6 +722,7 @@ mod canonical_path_tests { "pipeline", "scheduler", "sm100", + "sm100_attention", "tensor", "tiled_copy", ] @@ -1581,6 +1643,7 @@ pub(crate) fn try_translate_cute_call( } Some(match cute_fn { + CuteFn::Attention(kind) => sm100_attention_emit::emit(ctx, body, func, args, kind, target, block_ptr, prev_op, value_map, block_map, loc.clone()), CuteFn::Sm100(kind) => sm100_emit::emit( ctx, body, func, args, kind, target, block_ptr, prev_op, value_map, block_map, loc.clone(), ), @@ -2084,6 +2147,9 @@ pub(crate) fn try_translate_cute_call( 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, diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/sm100_attention_emit.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/sm100_attention_emit.rs new file mode 100644 index 0000000000..7811106c57 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/sm100_attention_emit.rs @@ -0,0 +1,217 @@ +/* + * 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::{Located, 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}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AttentionFn { + MmaSs, + MmaTs, + TmaLoad, + TmaStore, +} +fn invalid(message: impl core::fmt::Display) -> pliron::result::Error { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid SM100 CuTe boundary: {message}" + ))) +} + +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) +} + +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")) +} + +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) +} +fn shared_operands(kind: AttentionFn) -> &'static [usize] { + match kind { + AttentionFn::MmaSs => &[1, 2], + AttentionFn::MmaTs => &[2], + AttentionFn::TmaLoad => &[1, 5], + AttentionFn::TmaStore => &[1], + } +} +fn operand_count(kind: AttentionFn) -> usize { + match kind { + AttentionFn::MmaSs => 6, + AttentionFn::MmaTs => 5, + AttentionFn::TmaLoad => 7, + AttentionFn::TmaStore => 5, + } +} +#[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/static_config.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/static_config.rs index d8efd348e0..cbd1d0f437 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/static_config.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/static_config.rs @@ -300,6 +300,32 @@ fn decode_swizzle(ty: &Ty) -> Result { }) } +/// Decode the sealed attention layout identity and its major mode. +pub(crate) fn decode_attention_layout(ty: &Ty) -> Result<(ComposedLayout, bool), String> { + const PATH: &str = "cute_rs::sm100_attention::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, + )) +} + pub(crate) fn decode_smem_layout(ty: &Ty) -> Result { if let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() && canonical_adt_path(&def) == COMPOSED 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 1f776ed34a..3a4571b532 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/active_mask.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/active_mask.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `active_mask` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/clc.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/clc.rs index 79a46eece4..da8afed6e0 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/clc.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/clc.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `clc` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_barrier.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_barrier.rs index ea19c42e82..88c358c3b4 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_barrier.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `cluster_barrier` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_memory.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_memory.rs index 08785ee111..bd5ef35684 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_memory.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_memory.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `cluster_memory` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cp_async.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cp_async.rs index de9417e754..4e9da9a43d 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cp_async.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cp_async.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `cp_async` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/debug_control.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/debug_control.rs index 0a4f7d3cf8..9a5af42a11 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/debug_control.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/debug_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `debug_control` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/dotprod.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/dotprod.rs index b82d5a7a1f..44ace5f8f4 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/dotprod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/dotprod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `dotprod` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/elect.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/elect.rs index 32615073df..fa777d397f 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/elect.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/elect.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `elect` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/execution_control.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/execution_control.rs index 5bd8d1f057..edb3f04e63 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/execution_control.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/execution_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `execution_control` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/extended_minmax.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/extended_minmax.rs index 962e0d83c2..5fa87a6085 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/extended_minmax.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/extended_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `extended_minmax` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/helpers.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/helpers.rs index ce27000088..6a37f35fd7 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/helpers.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/helpers.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Shared glue for the generated intrinsic dispatch shards. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/integer_minmax.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/integer_minmax.rs index 55ed9eb623..b942cedd3e 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/integer_minmax.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/integer_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `integer_minmax` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/ldmatrix.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/ldmatrix.rs index f0ecafb79c..03ae6ad06c 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/ldmatrix.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/ldmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `ldmatrix` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_basic.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_basic.rs index ef3bb24e4a..80fdd89eb5 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_basic.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_basic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `mbarrier_basic` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_extended.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_extended.rs index d1a3f9d1b4..1ba1ce3ffb 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_extended.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_extended.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `mbarrier_extended` intrinsics. @@ -13,10 +13,12 @@ use super::helpers::require_arity; use crate::error::{TranslationErr, TranslationResult}; use crate::translator::{rvalue, terminator::helpers, values::ValueMap}; use dialect_nvvm::ops::{ - FenceMbarrierInitReleaseClusterOp, FenceProxyAsyncGenericAcquireSharedClusterClusterOp, + FenceAcquireSharedClusterClusterOp, FenceMbarrierInitReleaseClusterOp, + FenceProxyAsyncGenericAcquireSharedClusterClusterOp, FenceProxyAsyncGenericReleaseSharedCtaClusterOp, FenceProxyAsyncSharedCtaOp, - MbarrierArriveClusterOp, MbarrierArriveExpectTxClusterOp, MbarrierArriveExpectTxSharedOp, - MbarrierTryWaitParityClusterOp, MbarrierTryWaitParitySharedOp, MbarrierTryWaitSharedOp, + MbarrierArriveClusterOp, MbarrierArriveClusterRelaxedOp, MbarrierArriveExpectTxClusterOp, + MbarrierArriveExpectTxSharedOp, MbarrierTryWaitParityClusterOp, + MbarrierTryWaitParityClusterRelaxedOp, MbarrierTryWaitParitySharedOp, MbarrierTryWaitSharedOp, NanosleepOp, }; use pliron::basic_block::BasicBlock; @@ -41,6 +43,31 @@ pub(super) fn try_dispatch( loc: Location, ) -> TranslationResult>> { match name { + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033" + | "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:i1033"); + 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::i1031" + | "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:i1031"); + 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::i1032" + | "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: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, 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 f2fc0eb469..a5f38ee3c6 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility path dispatch for CUDA intrinsics. @@ -601,6 +601,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::i1033" + | "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 +832,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::i1031" + | "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 +852,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::i1032" + | "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" @@ -2447,6 +2453,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::i1033" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0313" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0315" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0314" @@ -2561,6 +2568,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::i1031" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0306" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0307" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1017" @@ -2570,6 +2578,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::i1032" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0838" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0067" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0830" @@ -3724,6 +3733,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::i1033" + | "cuda_device::barrier::fence_acquire_shared_cluster_cluster" => Some("v1:i1033"), "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" @@ -3960,6 +3971,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::i1031" + | "cuda_device::barrier::mbarrier_arrive_cluster_relaxed" => Some("v1:i1031"), "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" @@ -3978,6 +3991,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::i1032" + | "cuda_device::barrier::mbarrier_try_wait_parity_cluster_relaxed" => Some("v1:i1032"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0838" | "cuda_device::bf16::min_bf16" => { Some("v1:i0838") } 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 e43b58af87..260a42244c 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/movmatrix.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `movmatrix` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_alu.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_alu.rs index 489f14bc7c..6539f408e3 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_alu.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_alu.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `packed_alu` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_atomic.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_atomic.rs index 0340de31ec..fbbcb161f1 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_atomic.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_atomic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `packed_atomic` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_conversion.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_conversion.rs index 3e24585fa9..2880228a99 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_conversion.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `packed_conversion` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/prmt.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/prmt.rs index 1ac724c1c5..caaf7b7406 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/prmt.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `prmt` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/redux.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/redux.rs index 33cd317ebc..7f3d4ada4a 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/redux.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/redux.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `redux` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/register_mma.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/register_mma.rs index 2fdbfb149a..173ff7b202 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/register_mma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/register_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `register_mma` intrinsics. 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 6cbbc3c6f0..e0e7c4c96f 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_arithmetic.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_arithmetic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `scalar_arithmetic` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_conversion.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_conversion.rs index 1e5374feaa..d13cf29c95 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_conversion.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `scalar_conversion` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_math.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_math.rs index 24d4bf691a..bb8250d895 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_math.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_math.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `scalar_math` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sparse_mma.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sparse_mma.rs index 176166137c..079e886277 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sparse_mma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `sparse_mma` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sreg.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sreg.rs index 0a0805cc4d..095c0fa897 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sreg.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 2b3292696e..8dfeed0739 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/stmatrix.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `stmatrix` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sync.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sync.rs index 10649076e3..3fcade90eb 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sync.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sync.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `sync` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_cp.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_cp.rs index f1fc1d1dd5..a324314f77 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_cp.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_cp.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `tcgen05_cp` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_ld.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_ld.rs index 1b15405012..d15b928f64 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_ld.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_ld.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `tcgen05_ld` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_mma.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_mma.rs index 3a883ded27..ef42b3daba 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_mma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `tcgen05_mma` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_other.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_other.rs index fafa94eff7..b43c711446 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_other.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_other.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `tcgen05_other` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_st.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_st.rs index a5a4d2916a..24509559de 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_st.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_st.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `tcgen05_st` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/active_mask.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/active_mask.rs index 94ba989634..2b549aae5f 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/active_mask.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/active_mask.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 3adbcf384d..3a1856fe0a 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/clc.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/clc.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 6b79a6f627..386d567481 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cluster_barrier.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cluster_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 4cd7d0e9c7..3fd74cc440 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cluster_memory.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cluster_memory.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 732db6acc7..8a0eff3fef 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cp_async.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cp_async.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 aac10a5367..0ad2ac4ed5 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/debug_control.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/debug_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 4f2f47b6ba..f75d5e27bb 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/dotprod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/dotprod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 661867cb56..b27a63fe88 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/elect.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/elect.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 5982379821..9a5592e982 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/execution_control.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/execution_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 e4d1674115..39209587bf 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/extended_minmax.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/extended_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 743842a40b..a659221f09 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/integer_minmax.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/integer_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 c42b5322cb..bf3c458752 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/ldmatrix.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/ldmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 49cd96af2e..72ba281d43 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mbarrier_basic.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mbarrier_basic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 a485e9bcd2..cf5c644643 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: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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::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::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:i1033") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1033" + )); +} #[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::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::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:i1031") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1031" + )); +} +#[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::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::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:i1032") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1032" + )); +} +#[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 8e83cccd96..c01101e1e2 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 aeebe5fceb..58198d5063 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/movmatrix.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 b3b9d6f150..88d97ec10d 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_alu.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_alu.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 6c4946c578..4697a451b0 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_atomic.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_atomic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 3db97072cd..0239b8fbbf 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_conversion.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `packed_conversion` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/prmt.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/prmt.rs index 595d3cfb33..79086ac6c9 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/prmt.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 980f9627a3..9508dfee33 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/redux.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/redux.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 0313f3d356..e53c1c8b6e 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/register_mma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/register_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `register_mma` intrinsics. 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 28e9c45197..b02079b23d 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_arithmetic.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_arithmetic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 45bf50c365..d5186ffc47 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_conversion.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 261dfb42a5..a4c020c922 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_math.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_math.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 e91e9e25c5..92aafd7ec3 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sparse_mma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 0ba55dc947..96a3730709 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sreg.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 1903f837aa..765b0e569a 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/stmatrix.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 72099f07aa..86073be0dc 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sync.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sync.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 5fd2301bbf..c01775148f 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/tcgen05.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/tcgen05.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 c8a5313e42..b207b1df75 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/tma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `tma` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/vote.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/vote.rs index 68529a975b..989dda97b1 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/vote.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/vote.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 db7f3e09a9..1cb91d4755 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_barrier.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 1d968e101b..3e288d6ff1 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_match.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_match.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 24114fa6ec..5ba4855c58 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_shuffle.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_shuffle.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 cce6b57e33..2ddf80136b 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/wgmma_control.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 d957a4ff5d..8d56a043d6 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `tma` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/vote.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/vote.rs index 296b2fa571..a7f939a6dd 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/vote.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/vote.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `vote` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_barrier.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_barrier.rs index 9744f6f842..629b895a70 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_barrier.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `warp_barrier` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_match.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_match.rs index 164ad55bc1..5cefd50101 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_match.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_match.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `warp_match` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_shuffle.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_shuffle.rs index 1ee4c41b26..0a65445797 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_shuffle.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_shuffle.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `warp_shuffle` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/wgmma_control.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/wgmma_control.rs index 7892f9761f..23ab1d2ba4 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/wgmma_control.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `wgmma_control` intrinsics. 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 4902728862..c75560e061 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/active_mask.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/active_mask.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 3ed559ddd4..5e5c0acd3b 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/clc.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/clc.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 c4077828c3..30308c637c 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_barrier.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `cluster_barrier` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_memory.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_memory.rs index 99b41b142d..97a7d1851d 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_memory.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_memory.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 e3be86302f..f508e95e92 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cp_async.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cp_async.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 c113446cd1..51f117eb46 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/debug_control.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/debug_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `debug_control` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/dotprod.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/dotprod.rs index 517873c6c9..46457534dc 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/dotprod.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/dotprod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 8675140556..940265effb 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/elect.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/elect.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 c87931d904..4d05a3955e 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/execution_control.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/execution_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 40905729d4..10ec1d67ca 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/extended_minmax.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/extended_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `extended_minmax` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/integer_minmax.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/integer_minmax.rs index 623224bea5..a300244dce 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/integer_minmax.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/integer_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 664b89a0c6..e590e788bb 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/ldmatrix.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/ldmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `ldmatrix` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_basic.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_basic.rs index ae6461e34d..b16fb697ca 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_basic.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_basic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 5e9345d9a9..b8f89b7b05 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_extended.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_extended.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 9d75fd2229..93ada91d41 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mod.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 6df686bace..b868f92356 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/movmatrix.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 0679dbe97d..a1b61765eb 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_alu.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_alu.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 72316705f0..7fe8e1470d 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_atomic.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_atomic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `packed_atomic` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_conversion.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_conversion.rs index 5bf7f6a9e6..058491b1cc 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_conversion.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `packed_conversion` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/prmt.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/prmt.rs index b7478d4c80..7d826d5402 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/prmt.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `prmt` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/redux.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/redux.rs index 6996ba8286..358088cc1d 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/redux.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/redux.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 d03e45f80a..d9293a7fd5 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/register_mma.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/register_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `register_mma` intrinsics. 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 34f15ebb7a..edb1f2c68a 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_arithmetic.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_arithmetic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `scalar_arithmetic` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_conversion.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_conversion.rs index 0143974701..ef683f04c9 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_conversion.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `scalar_conversion` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_math.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_math.rs index 1cfada322c..798d3c17e0 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_math.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_math.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `scalar_math` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sparse_mma.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sparse_mma.rs index bafded4d4a..34cc80fa64 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sparse_mma.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `sparse_mma` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sreg.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sreg.rs index 6924af4a02..8b9c6ee94e 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sreg.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 45322f236d..5aa0429e5a 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/stmatrix.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 27d950019a..b35399e941 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sync.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sync.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 87bd3bed28..3bcfa5b3c6 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tcgen05.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tcgen05.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `tcgen05` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tma.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tma.rs index 1f4a6c8608..9b3d063854 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tma.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `tma` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/vote.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/vote.rs index 5f1507ffa4..99df456556 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/vote.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/vote.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 0c6704bb83..22606d8804 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_barrier.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 059548262b..26a84712b7 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_match.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_match.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 dd32411eed..3e699a0edf 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_shuffle.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_shuffle.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // 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 18d3f097ae..09ad6a6527 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/wgmma_control.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `wgmma_control` intrinsics. 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 622fbb0ac5..f93abfb3c4 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 @@ -1191,8 +1191,8 @@ mod tests { .filter(|&o| Operation::get_op::(o, &ctx).is_none()) .collect(); for op in moved { - op.unlink(&mut ctx); - op.insert_before(&mut ctx, if before { store } else { goto }); + 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 @@ -1209,7 +1209,7 @@ mod tests { vec![], 0, ); - use_value.insert_after(&mut ctx, load); + use_value.insert_after(&ctx, load); if escape { let pointer = fixture.alloca.deref(&ctx).get_result(0); let ty = pointer.get_type(&ctx); @@ -1222,7 +1222,7 @@ mod tests { 0, ); MirCastOp::new(cast).set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - cast.insert_before(&mut ctx, goto); + cast.insert_before(&ctx, goto); } canonicalize_read_only_aggregate_arguments(fixture.module, &mut ctx, false); assert_eq!( 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..5bfe42f497 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..756d456a5b 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..3411d403e9 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0707f363f0 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..c226399704 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..b62a72b4a1 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..150c36e13b 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..f861aa38b9 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..bcb0e718a4 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..86fe4bf443 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 @@ -204,6 +204,7 @@ dependencies = [ "cuda-artifact-finalizer", "cuda-async", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..ebd4364b37 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 @@ -204,6 +204,7 @@ dependencies = [ "cuda-artifact-finalizer", "cuda-async", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..33e6b20a77 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..4d735e003b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/Cargo.lock @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..5df295622d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/Cargo.lock @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..7342cbacaa 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..ca223c09e4 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..dd7ef5e9fb 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..9a1ec5e3ad 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..4df70740a3 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..cfa44698a5 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..f226544eba 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..958085f4ff 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..57e6dcba18 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/Cargo.lock @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..d1c32bb250 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/Cargo.lock @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..072aa71b48 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 @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..74be8b050d 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..97efed7d99 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..a791fda4dd 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..743d5e4286 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 @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..84398d774c 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..8de2f436db 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..bb76741cc7 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..6ff40dc52f 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..58b7409069 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..1b15f019a5 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..3bdec1e66d 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0c20922d16 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..2b839ae0c3 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..99efab64e3 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..4924f43949 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..03b2f871aa 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0dc6a2baf2 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 @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..eefd8ce1f1 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 @@ -179,6 +179,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0a91cb91fa 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 @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..64e2f273b3 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 @@ -182,6 +182,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..4a98df0527 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 @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..f2b2891b71 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 @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..71603ed322 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..a7f91e9631 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..7e27222404 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..b1c3278a2e 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..c92ccfea6e 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..e4d9b8f5f3 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..4ad3919cea 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..204abc15f1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/Cargo.lock @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..b26f0804e6 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..e9b923b382 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..a7a6977f55 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..128aac9e00 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..7456cfbe0f 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..3a8a6e75d7 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..e539a8a9c7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/Cargo.lock @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..f8289c6bcb 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..3f42e0c743 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..3230197846 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..e36c03801e 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..26962bc4cb 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..a2773af37c 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..f0cd93e43c 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..933aa2bf66 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..2ee2097249 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..211e0f0608 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..7606450ac9 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..d5de59908e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error/Cargo.lock @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..55f803701c 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..9c1a9e7803 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..391e870c28 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..c9d53ab4de 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..36e9b9b97a 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..46c18664f3 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..c5d1e68fc3 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..e00d8de680 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..8e8eb2b74e 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..2a5043de98 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..9f53c12007 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..b5e78c6463 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..1f3ab91d9f 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..f29d481050 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..52b01c3a97 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..570cc4cc7b 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..829433bbce 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..6604248686 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..e70d10a7d0 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..cd70a9b87e 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..d00f959bfa 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..ae65d5ba00 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..308d0368a4 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..70799d7ca8 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..54db720c53 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..d26fe7b243 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..18ec08af14 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..cbbab605d4 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..e03e9e7d47 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..3fd26d9deb 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/Cargo.lock @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..3cf30c8391 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..88d559e1a1 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..37698dc184 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..8ed4fa9cc3 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..072d70fd4c 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..bd14dd242d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/Cargo.lock @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..5f18567469 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/Cargo.lock @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..fbed1d6a0c 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 @@ -202,6 +202,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..e7fd9ddee4 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 @@ -202,6 +202,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..c13ca9102b 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..9ce6f55f60 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..902b6541e3 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 @@ -193,6 +193,7 @@ dependencies = [ "cuda-artifact-finalizer", "cuda-async", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..75e7307cab 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..1310ee6db1 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..8f76362681 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 @@ -183,6 +183,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..440913683e 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..fd67415dfc 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..d4b9b5461c 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..56e29e736c 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0ebb8ed595 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..a8127f0e80 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..3679781adb 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 @@ -193,6 +193,7 @@ dependencies = [ "cuda-artifact-finalizer", "cuda-async", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..024e7021c8 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..61bbf580a1 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..ea948d4ecb 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..d02aed3412 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..842fe799c2 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..b9b4443170 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..b67b0c2820 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..f8ea8dac57 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0374d68288 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..57d3680b90 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..1e5c31789b 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..1006cc1840 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0a71113abd 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..a525a1063b 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 @@ -191,6 +191,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..738d09fad1 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..efff1f567e 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..4dc7e05f8b 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..a3e19f471c 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..a301de3fd7 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0c8304d674 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..41f152b232 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..7213d9a1d4 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..20f917aea5 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..e8131ada6c 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..fa6c46c499 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..fc8efbf260 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..5270ecf8c3 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0ff94aff67 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..403108d973 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..09eb25be73 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..9c368882b2 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..c2b35702a1 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..da7573a7f6 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..605b50cb6c 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..2115202334 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/Cargo.lock @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..d20300dad2 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..8808ce79c7 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..4cd3626406 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..05bdac2637 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..76d7a8c0a1 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..b69f61e9f3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/Cargo.lock @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0d204dee81 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..e6735f7e2e 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..e544562ba8 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..ca498ab154 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..f296378d74 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..67d816ee89 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..998043a456 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0c1512078b 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..3f7d11f6ac 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..4655916ada 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..b051e22b45 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..4078c13594 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..d430edba87 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..5e82b37ad7 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..9d6069aad4 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..78f0e3200f 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..4d29b3a3ad 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..cb93313c29 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..508c19dcaa 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..060c6ebabc 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/Cargo.lock @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..1c0e33cde5 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..269e985a33 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..9a19381a14 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..9482c9c8e9 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..807aab7821 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0f86ff4720 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..2a9dbb0190 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..d1038b41ab 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..6d493bb555 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..2c77b40a20 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..ca75ac4f4c 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..3e900ce907 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0fc9230fb1 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..34652fd481 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..fac95ab46d 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..74e32551dc 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..59ab7c0970 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..93375fc554 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..751ca7ee27 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..6610665357 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/Cargo.lock @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..b030857670 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..f578ba62cb 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..851d58560b 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..ed4e4630c2 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..7bf24c4da1 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..844ce99fc4 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..8c0b170c27 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..6153e10851 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..04aa6fb4cd 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..38ec439a74 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..805711f85b 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..2c7e027be5 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..68f0751f0e 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..0d509e67ec 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..a2f2d7f29c 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..93ce620ff5 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..154c0b359c 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..2a4a9c7751 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/Cargo.lock @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..c634d101a5 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..53b2b2831f 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..d56f2593c8 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..d657d9fd8d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/Cargo.lock @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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..aff2de7fc9 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 @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", 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 8c326a557b..b8ae4db197 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/src/generated_intrinsics.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/src/generated_intrinsics.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated collector predicates for intrinsic placeholder functions. @@ -243,6 +243,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::i1033", "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 +358,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::i1031", "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 +368,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::i1032", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0838", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0067", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0830", @@ -1276,6 +1279,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", @@ -1390,6 +1394,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", @@ -1399,6 +1404,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", @@ -2315,6 +2321,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::i1033" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0313" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0315" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0314" @@ -2429,6 +2436,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::i1031" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0306" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0307" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1017" @@ -2438,6 +2446,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::i1032" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0838" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0067" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0830" @@ -3355,6 +3364,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" @@ -3470,6 +3480,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" @@ -3479,6 +3490,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" diff --git a/cuda-oxide/crates/rustc-codegen-cuda/src/lib.rs b/cuda-oxide/crates/rustc-codegen-cuda/src/lib.rs index 06801082a4..d1a8c1de35 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/src/lib.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/src/lib.rs @@ -287,6 +287,7 @@ //! | `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 | //! @@ -468,11 +469,28 @@ fn device_backend_from_env() -> Result { return Err("CUDA_OXIDE_MLIR_PROFILE is not valid Unicode".to_owned()); } }; - parse_device_backend( + 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( diff --git a/cuda-oxide/cuda-oxide-book/compiler/architecture-overview.md b/cuda-oxide/cuda-oxide-book/compiler/architecture-overview.md index f55ba48f53..8112bcbba2 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/architecture-overview.md +++ b/cuda-oxide/cuda-oxide-book/compiler/architecture-overview.md @@ -140,6 +140,7 @@ cuda-oxide is split into focused crates. Here is every one and its role: | `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 | diff --git a/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md b/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md index 8cff3c042a..464f6a9c5e 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md +++ b/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md @@ -303,16 +303,16 @@ they become `call` instructions to `@llvm.nvvm.*` intrinsics. ### Architecture Coverage -At catalog SHA-256 `c243e568` (the stamp in every `ops/generated/` file -header), the dialect holds 576 operations across 42 modules, and they come +At catalog SHA-256 `b8c889f1` (the stamp in every `ops/generated/` file +header), the dialect holds 580 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 `c243e568`, the counts on this page predate the +stamp no longer starts with `b8c889f1`, the counts on this page predate the catalog you are reading. **Hand-written**, directly under `cuda-oxide/crates/dialect-nvvm/src/ops/`. These are the ops with bespoke verification or lowering that the intrinsic catalog does not -describe. There are seven modules and 26 operations: +describe. There are seven modules and 27 operations: | Module | Description | Ops | | :-------- | :---------------------------------------------------------- | --: | @@ -321,13 +321,13 @@ describe. There are seven modules and 26 operations: | `cluster` | Cluster index and cluster-count registers | 2 | | `debug` | `assertfail`, `vprintf` | 2 | | `grid` | Cooperative `grid_sync` | 1 | -| `memory` | Generic-to-shared address conversion with a byte offset | 1 | +| `memory` | Generic-to-shared address conversion and global stores | 2 | | `wgmma` | Warpgroup MMA descriptors; bf16/f16 at m64n64k16, bf16 at m64n128k16, tf32 at m64n64k8 | 14 | **Generated**, under `ops/generated/`, from `intrinsics/catalog.json` by `cuda-intrinsics-gen`. Every file there opens with `// @generated ... DO NOT EDIT.`, and editing one by hand is undone by the next generator run. This is -the large majority -- 35 modules and 550 operations, resolved from 1029 catalog +the large majority -- 35 modules and 553 operations, resolved from 1033 catalog entries, since several intrinsics can share one structural op: | Area | Modules | Ops | @@ -336,7 +336,7 @@ entries, since several intrinsics can share one structural op: | Tensor Memory Accelerator | `tma` | 111 | | 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 | +| 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/cuteir/README.md b/cuda-oxide/cuteir/README.md index ca21d9cd56..8575fc1373 100644 --- a/cuda-oxide/cuteir/README.md +++ b/cuda-oxide/cuteir/README.md @@ -63,7 +63,21 @@ lifetimes remain the caller's responsibility. | [`fp16_gemm_256x352_cute`](examples/fp16_gemm_256x352_cute) | SM100 two-CTA GEMM with typed shared tiles, TMEM operations, collector reuse, and cluster TMA pipelines through the CUTLASS translation backend. | | [`fp16_gemm_256x352_ctm`](examples/fp16_gemm_256x352_ctm) | The same SM100 FP16 GEMM through the typed `ctm-rs` instruction API, with explicit TMA stages, MMA issue, and TMEM completion. | -The examples are deliberately ordered from smallest to largest: +The datacenter Blackwell attention examples target **SM100a/B200**: + +| Example | Matched CuTe-RS and pinned CuTeDSL implementations | +| :--- | :--- | +| [`sm100_fmha_prefill_cute`](examples/sm100_fmha_prefill_cute) | FP16 D128 prefill, dense or causal MHA/GQA, with FP32 accumulation and persistent warp specialization. | +| [`sm100_mla_decode_fp16_cute`](examples/sm100_mla_decode_fp16_cute) | FP16 MLA decode with a paged latent cache, two-CTA clusters, and optional split-KV reduction. | + +Each example vendors its CUTLASS source and helpers under `reference/`, checks +both outputs against an independent oracle, and compares complete invocations +with the same event-timing policy. Kernel names distinguish `_cute_rs` and +`_cutedsl`. See the example READMEs for supported shape profiles and validation +commands. These use CuTe layouts plus NVVM operations; `cute_ext` GQA examples +belong to the separate LIR programming model. + +The original examples are deliberately ordered from smallest to largest: ```text elementwise GEMV GEMM diff --git a/cuda-oxide/cuteir/ctm-rs/src/cluster.rs b/cuda-oxide/cuteir/ctm-rs/src/cluster.rs index 09d79eb605..9c992354d7 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/cluster.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/cluster.rs @@ -161,3 +161,23 @@ pub unsafe fn arrive_remote(address: ClusterSharedAddress<'_, Barrier>) { 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/cuteir/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.rs new file mode 100644 index 0000000000..5b09b88092 --- /dev/null +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.rs @@ -0,0 +1,10 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use ctm_rs::{cluster::arrive_remote_relaxed, memory::SharedAddress}; +use cuda_device::barrier::Barrier; +fn wrong(local: SharedAddress) { + unsafe { arrive_remote_relaxed(local) }; +} +fn main() {} diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.stderr new file mode 100644 index 0000000000..31ecb1b721 --- /dev/null +++ b/cuda-oxide/cuteir/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:8:36 + | +8 | 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/cuteir/ctm-rs/tests/primitives_types.rs b/cuda-oxide/cuteir/ctm-rs/tests/primitives_types.rs index 8404052019..6cafc9a8c4 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/primitives_types.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/primitives_types.rs @@ -57,3 +57,10 @@ fn multicast_masks_and_signed_coordinates_preserve_instruction_domains() { assert_eq!(Coordinates2D::new(-32, 4096).inner(), -32); assert_eq!(Coordinates2D::new(-32, 4096).outer(), 4096); } + +#[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; +} diff --git a/cuda-oxide/cuteir/cute-rs/Cargo.toml b/cuda-oxide/cuteir/cute-rs/Cargo.toml index eec6a8a1d8..466425a575 100644 --- a/cuda-oxide/cuteir/cute-rs/Cargo.toml +++ b/cuda-oxide/cuteir/cute-rs/Cargo.toml @@ -8,6 +8,7 @@ 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 } @@ -17,4 +18,6 @@ cuda-core = { workspace = true, optional = true } 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/cuteir/cute-rs/examples/global_store_compile.rs b/cuda-oxide/cuteir/cute-rs/examples/global_store_compile.rs new file mode 100644 index 0000000000..5b1ffd74f2 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/examples/global_store_compile.rs @@ -0,0 +1,62 @@ +// 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); + } +} + +fn main() {} // Kernels are compiled, never launched. diff --git a/cuda-oxide/cuteir/cute-rs/src/arch.rs b/cuda-oxide/cuteir/cute-rs/src/arch.rs new file mode 100644 index 0000000000..9ea21cc2ac --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/arch.rs @@ -0,0 +1,23 @@ +/* + * 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. + +/// Return a warp-uniform value in a form that can favor uniform-register use. +/// +/// This matches CuTe's `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 assume_warp_uniform(value: u32) -> u32 { + cuda_device::warp::shuffle_sync(u32::MAX, value, 0) +} diff --git a/cuda-oxide/cuteir/cute-rs/src/lib.rs b/cuda-oxide/cuteir/cute-rs/src/lib.rs index 97df4ee7fa..60ddf54b8f 100644 --- a/cuda-oxide/cuteir/cute-rs/src/lib.rs +++ b/cuda-oxide/cuteir/cute-rs/src/lib.rs @@ -57,6 +57,7 @@ #[cfg(feature = "host")] extern crate alloc; +pub mod arch; pub mod block_scaled; pub mod block_scaled_mma; pub mod cooperative; @@ -67,6 +68,7 @@ pub mod numeric; pub mod pipeline; pub mod scheduler; pub mod sm100; +pub mod sm100_attention; pub mod tensor; pub mod tile; pub mod tiled_copy; @@ -102,6 +104,9 @@ pub use sm100::{ Sm100TiledMma, Sm100TmaMmaPipeline, Sm100TmaStorePipeline, Sm100Tmem, Sm100TmemEpilogue, }; pub use tensor::{Contiguous1D, RegisterTile, Tensor, TensorElement, TensorMut, Tile1D, Zipped1D}; -pub use tile::{assume_div, load_tile, store_tile}; +pub use tile::{ + GlobalStoreElement, GlobalStoreTile, assume_div, load_tile, store_global_no_allocate, + store_tile, +}; pub use tiled_copy::{GlobalCopyTensor, SharedTensor, TiledCopy}; pub use tma::{TmaDesc, copy_tma_2d, copy_tma_s2g_2d, tile_bytes}; diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention.rs new file mode 100644 index 0000000000..83ecacff87 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention.rs @@ -0,0 +1,536 @@ +/* + * 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. +pub mod pipeline; + +use crate::{OperandA, OperandB, TmaDesc}; +use core::marker::PhantomData; +use cuda_device::{Barrier, CuSimd}; + +/// Physical layout of a shared FP16 operand, with CuTe's tiled byte swizzle. +mod sealed { + 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. + fn offset(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. +pub struct SharedLayout; +impl SharedLayout { + 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. + #[inline(always)] + pub fn offset(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 offset(row: u32, col: u32) -> usize { + Self::offset(row, col) + } +} + +/// Non-owning shared tile, typed by its physical layout and MMA operand role. +pub struct SharedTile { + base: *mut f16, + byte_offset: u32, + 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 { + base: *mut f16, + marker: PhantomData<(L, Role)>, +} +impl SharedTileArray { + 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 + }; + 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, + } + } +} + +/// Independent shared/shared FP16 MMA with FP32 accumulation. +pub struct MmaSs( + PhantomData<(AL, BL)>, +); +impl + MmaSs +{ + 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. +/// The importer currently supports K=32 or a positive multiple of 64. +pub struct MmaTs(PhantomData); +impl MmaTs { + 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() + } +} + +/// 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], +} +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) }; + cuda_device::tcgen05::tcgen05_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 { cuda_device::tcgen05::$load(address) }.to_array(); + let mut values = [0.0; $n]; + let mut i = 0; + while i < $n { + cuda_device::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) }; + cuda_device::tcgen05::tcgen05_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 { + cuda_device::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) }; + cuda_device::tcgen05::tcgen05_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 { cuda_device::tcgen05::$store(address, CuSimd::new(raw)) }; + } + } + }; +} +fragment!(2, tcgen05_ld_32x32b_x2_raw, tcgen05_st_32x32b_x2_raw); +fragment!(4, tcgen05_ld_32x32b_x4_raw, tcgen05_st_32x32b_x4_raw); +fragment!(8, tcgen05_ld_32x32b_x8_raw, tcgen05_st_32x32b_x8_raw); +fragment!(16, tcgen05_ld_32x32b_x16_raw, tcgen05_st_32x32b_x16_raw); +fragment!(32, tcgen05_ld_32x32b_x32_raw, tcgen05_st_32x32b_x32_raw); +// 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. +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) }; + cuda_device::tcgen05::tcgen05_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 { cuda_device::tcgen05::tcgen05_ld_32x32b_x32_raw(address + $offset) };)+ + let mut values = [0.0; $n]; + $( + let raw = $chunk.to_array(); + let mut i = 0; + while i < 32 { + cuda_device::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) }; + cuda_device::tcgen05::tcgen05_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 { + cuda_device::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) }; + cuda_device::tcgen05::tcgen05_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 { + cuda_device::thread::__unroll_config::<0>(); + chunk[i] = raw[$offset + i]; + i += 1; + } + unsafe { cuda_device::tcgen05::tcgen05_st_32x32b_x32_raw(address + $offset, CuSimd::new(chunk)) }; + )+ + } + } + } +} +chunked_fragment!(64; 0 => a, 32 => b); +chunked_fragment!(128; 0 => a, 32 => b, 64 => c, 96 => d); + +/// A rank-three TMA copy into the same physical layout used by MMA. +pub struct TmaCopy(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: &TmaDesc, + 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: &TmaDesc, 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, + ) + } + } +} + +#[doc(hidden)] +pub mod __compiler { + use super::*; + #[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") + } + #[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") + } + #[inline(never)] + #[allow(clippy::too_many_arguments)] + pub unsafe fn tma_load( + desc: *const TmaDesc, + 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") + } + #[inline(never)] + pub unsafe fn tma_store( + desc: *const TmaDesc, + src: *mut f16, + c0: i32, + c1: i32, + c2: i32, + ) { + let _ = (desc, src, c0, c1, c2); + unreachable!("CuTe compiler marker") + } +} + +#[cfg(feature = "host")] +mod host; +#[cfg(feature = "host")] +pub use host::make_tma_desc; diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/host.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/host.rs new file mode 100644 index 0000000000..be592f9209 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/host.rs @@ -0,0 +1,92 @@ +/* + * 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::TmaDesc; +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 matrix into the selected shared operand layout. +/// +/// Large contiguous dimensions are split into swizzle-sized chunks. For an +/// MN-major operand the logical rows/columns are transposed relative to global +/// physical rows/columns. Global coordinates always use the physical ordering. +/// The descriptor borrows the allocation; callers retain it until completion. +#[allow(clippy::not_unsafe_ptr_arg_deref)] +pub fn make_tma_desc( + 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) % 16 != 0 { + 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(TmaDesc { + bytes: unsafe { core::mem::transmute::(out.assume_init()) }, + marker: PhantomData, + }) +} diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/pipeline.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/pipeline.rs new file mode 100644 index 0000000000..5a219f5299 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/pipeline.rs @@ -0,0 +1,584 @@ +/* + * 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. + +use core::marker::PhantomData; +use core::mem::{align_of, size_of}; +use ctm_rs::{cluster, groups::CtaGroup2, memory::SharedAddress, tcgen05}; +use cuda_device::barrier::{self, Barrier}; + +mod sealed { + pub trait Wait { + // Existing policies order a subsequent TC receiver. Only explicitly + // non-TC receiver policies may omit this fence. + const TC_AFTER: bool = true; + } + pub trait Signal {} + pub trait Mma {} + pub trait Tma {} + 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 { barrier::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 { barrier::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 { barrier::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 { barrier::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 { barrier::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 { barrier::mbarrier_try_wait_parity_cluster_relaxed(bar, phase) } {} + unsafe { barrier::fence_proxy_async_generic_acquire_shared_cluster_cluster() }; + } +} +impl SignalPolicy for SharedToAsync { + #[inline(always)] + unsafe fn arrive(bar: *mut Barrier) { + unsafe { + barrier::fence_proxy_async_generic_release_shared_cta_cluster(); + 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 { barrier::mbarrier_try_wait_parity_cluster_relaxed(bar, phase) } {} + unsafe { barrier::fence_acquire_shared_cluster_cluster() }; + } +} + +/// Producer position: initially waits for empty phase one. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct Producer; +/// Consumer position: initially waits for full phase zero. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct Consumer; +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 { + #[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. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct RingCursor { + index: u32, + phase: u32, + role: PhantomData, +} +impl RingCursor { + 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. + #[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. +#[derive(Clone, Copy)] +pub struct TwoCtaRing { + full: *mut Barrier, + empty: *mut Barrier, + ordering: PhantomData<(Full, Empty)>, +} + +impl TwoCtaRing { + 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)] + const fn slot(ticket: u32) -> usize { + (ticket % STAGES) as usize + } + #[inline(always)] + const fn full_phase(ticket: u32) -> u32 { + (ticket / STAGES) & 1 + } + #[inline(always)] + 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 { + barrier::mbarrier_init(self.full.add(stage as usize), producers); + barrier::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 expect(self, ticket: u32, bytes: u32) + where + Full: TmaPolicy, + { + unsafe { barrier::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::expect`] apply to the current cursor position. + #[inline(always)] + pub unsafe fn expect_cursor(self, cursor: RingCursor, bytes: u32) + where + Full: TmaPolicy, + { + unsafe { + barrier::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)] +fn after_wait() { + if ::TC_AFTER { + tcgen05::fence_after_thread_sync(); + } +} + +#[inline(always)] +unsafe fn arrive_remote(bar: *mut Barrier) { + unsafe { + cluster::arrive_remote_relaxed(cluster::map_shared(SharedAddress::from_raw(bar), 0)); + } +} +#[inline(always)] +unsafe fn mma_arrive(bar: *mut Barrier) { + unsafe { + tcgen05::commit_multicast_cg2( + SharedAddress::from_raw(bar), + cluster::CtaMask::::try_from_bits(3).unwrap_unchecked(), + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn tickets_wrap_slots_and_flip_opposite_phases() { + 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); + } + } + #[test] + fn persistent_work_keeps_phase_at_a_partial_ring_boundary() { + 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); + } + } + #[test] + fn one_stage_ring_flips_phase_on_every_ticket() { + 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/cuteir/cute-rs/src/tile.rs b/cuda-oxide/cuteir/cute-rs/src/tile.rs index 081f8be0dd..8321838586 100644 --- a/cuda-oxide/cuteir/cute-rs/src/tile.rs +++ b/cuda-oxide/cuteir/cute-rs/src/tile.rs @@ -102,3 +102,57 @@ pub unsafe fn assume_div(x: usize) -> usize { // a different rule because every thread must reach nearby barriers: // // every thread calls copy_g2s -> its assigned part copies 16 / fewer / 0 bytes + +mod global_store_element { + 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(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/cuteir/cute-rs/src/tma.rs b/cuda-oxide/cuteir/cute-rs/src/tma.rs index 823560c31d..bd45d451fa 100644 --- a/cuda-oxide/cuteir/cute-rs/src/tma.rs +++ b/cuda-oxide/cuteir/cute-rs/src/tma.rs @@ -59,6 +59,23 @@ pub struct TmaDesc { pub marker: PhantomData<(T, SmemL)>, } +impl TmaDesc { + /// Prefetch this descriptor into the device's tensor-map cache. + /// + /// 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(&self) { + unsafe { ctm_rs::tma::prefetch_descriptor(core::ptr::from_ref(self).cast()) }; + } +} + // The type parameters describe the encoded bytes; they are not stored values. impl Copy for TmaDesc {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/arch.rs b/cuda-oxide/cuteir/cute-rs/tests/arch.rs new file mode 100644 index 0000000000..f39dd3a98d --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/arch.rs @@ -0,0 +1,10 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#[test] +fn warp_uniform_hint_requires_an_explicit_unsafe_u32_contract() { + let _: unsafe fn(u32) -> u32 = cute_rs::arch::assume_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"); +} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_storage.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_storage.rs new file mode 100644 index 0000000000..1d60cafe52 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_storage.rs @@ -0,0 +1,6 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +fn main() { + let _ = cute_rs::GlobalStoreTile([0.0f64; 4]); +} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_storage.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_storage.stderr new file mode 100644 index 0000000000..1f99fc57e3 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_storage.stderr @@ -0,0 +1,22 @@ +error[E0277]: the trait bound `f64: GlobalStoreElement` is not satisfied + --> tests/compile_fail/global_store_storage.rs:5:38 + | +5 | 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/tile.rs + | + | impl GlobalStoreElement for f16 {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `f16` + | impl GlobalStoreElement for f32 {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `f32` + | impl GlobalStoreElement for u32 {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `u32` +note: required by a bound in `GlobalStoreTile` + --> src/tile.rs + | + | pub struct GlobalStoreTile(pub [T; N]); + | ^^^^^^^^^^^^^^^^^^ required by this bound in `GlobalStoreTile` diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_unsafe.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_unsafe.rs new file mode 100644 index 0000000000..5655b0b9f8 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_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 values = cute_rs::GlobalStoreTile([0.0f32; 8]); + cute_rs::store_global_no_allocate(core::ptr::null_mut(), &values); +} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_unsafe.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_unsafe.stderr new file mode 100644 index 0000000000..4d524fedc4 --- /dev/null +++ b/cuda-oxide/cuteir/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:6:5 + | +6 | 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/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.rs new file mode 100644 index 0000000000..68ac5fc346 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.rs @@ -0,0 +1,6 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +use cute_rs::sm100_attention::{MmaSs, SharedLayout}; +type UnsupportedK16 = SharedLayout<128, 16>; +const _: MmaSs = MmaSs::new(); +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.stderr new file mode 100644 index 0000000000..aceb88a10d --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.stderr @@ -0,0 +1,26 @@ +error[E0080]: evaluation panicked: assertion failed: contiguous >= 32 && contiguous.is_multiple_of(32) + --> src/sm100_attention.rs + | + | impl SharedLayout { + | const VALID: () = { +... + | assert!(contiguous >= 32 && contiguous.is_multiple_of(32)); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::sm100_attention::SharedLayout::<128, 16>::VALID` failed here + +note: erroneous constant encountered + --> src/sm100_attention.rs + | + | const VALID: () = SharedLayout::::VALID; + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +note: erroneous constant encountered + --> src/sm100_attention.rs + | + | let () = AL::VALID; + | ^^^^^^^^^ + +note: erroneous constant encountered + --> src/sm100_attention.rs + | + | let () = Self::VALID; + | ^^^^^^^^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_operand_role.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_operand_role.rs new file mode 100644 index 0000000000..349fb8f34b --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_operand_role.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::sm100_attention::{MmaSs, SharedLayout, SharedTile}; +use cute_rs::OperandB; +type L = SharedLayout<128, 128>; +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/cuteir/cute-rs/tests/compile_fail/sm100_attention_operand_role.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_operand_role.stderr new file mode 100644 index 0000000000..d12a1611a1 --- /dev/null +++ b/cuda-oxide/cuteir/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:11:26 + | +11 | 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/sm100_attention.rs + | + | pub unsafe fn gemm( + | ^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.rs new file mode 100644 index 0000000000..35b104bc50 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.rs @@ -0,0 +1,13 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#![feature(f16)] +use cute_rs::OperandB; +use cute_rs::sm100_attention::{SharedLayout, SharedTileArray}; +type L = SharedLayout<8, 32>; +const EMPTY: SharedTileArray = + unsafe { SharedTileArray::from_raw(core::ptr::null_mut()) }; +const OVERFLOW: SharedTileArray = + unsafe { SharedTileArray::from_raw(core::ptr::null_mut()) }; +fn main() { + let _ = (EMPTY, OVERFLOW); +} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.stderr new file mode 100644 index 0000000000..1a4ab1eeb6 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.stderr @@ -0,0 +1,24 @@ +error[E0080]: evaluation panicked: assertion failed: TILES > 0 + --> src/sm100_attention.rs + | + | impl SharedTileArray { +... + | const VALID: () = { + | assert!(TILES > 0); + | ^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::sm100_attention::SharedTileArray::, cute_rs::OperandB, 0>::VALID` failed here + +note: erroneous constant encountered + --> src/sm100_attention.rs + | + | let () = Self::VALID; + | ^^^^^^^^^^^ + +error[E0080]: evaluation panicked: assertion failed: TILES <= u32::MAX / Self::TILE_BYTES + --> src/sm100_attention.rs + | + | impl SharedTileArray { +... + | const VALID: () = { + | assert!(TILES > 0); + | assert!(TILES <= u32::MAX / Self::TILE_BYTES); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::sm100_attention::SharedTileArray::, cute_rs::OperandB, u32::MAX>::VALID` failed here diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.rs new file mode 100644 index 0000000000..2b65e4d3c3 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.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::OperandB; +use cute_rs::sm100_attention::{MmaSs, SharedLayout, SharedTileArray}; +type L = SharedLayout<64, 64>; +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/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.stderr new file mode 100644 index 0000000000..4ecd0009b5 --- /dev/null +++ b/cuda-oxide/cuteir/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:11:26 + | +11 | 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/sm100_attention.rs + | + | pub unsafe fn gemm( + | ^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tma_layout.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tma_layout.rs new file mode 100644 index 0000000000..00cf1f779f --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tma_layout.rs @@ -0,0 +1,11 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#![feature(f16)] +use cute_rs::sm100_attention::{SharedLayout, SharedTile, TmaCopy}; +use cute_rs::{OperandB, TmaDesc}; +type Qk = SharedLayout<128, 128>; +type Pv = SharedLayout<128, 128, true>; +fn mismatch(desc: &TmaDesc, tile: &mut SharedTile) { + unsafe { TmaCopy::::load(desc, tile, 0, 0, core::ptr::null_mut(), 0) }; +} +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tma_layout.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tma_layout.stderr new file mode 100644 index 0000000000..7ac84cc134 --- /dev/null +++ b/cuda-oxide/cuteir/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:9:43 + | +9 | 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/sm100_attention.rs + | + | pub unsafe fn load( + | ^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.rs new file mode 100644 index 0000000000..97d8eb31e1 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.rs @@ -0,0 +1,7 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +use cute_rs::sm100_attention::{MmaTs, SharedLayout}; +// MN-major shared B alone permits K16, but the TMEM-source importer does not. +type Value = SharedLayout<128, 16, true>; +const _: MmaTs = MmaTs::new(); +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.stderr new file mode 100644 index 0000000000..df70ab423b --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.stderr @@ -0,0 +1,14 @@ +error[E0080]: evaluation panicked: assertion failed: K == 32 || (K > 0 && K.is_multiple_of(64)) + --> src/sm100_attention.rs + | + | impl MmaTs { + | const VALID: () = { +... + | assert!(K == 32 || (K > 0 && K.is_multiple_of(64))); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::sm100_attention::MmaTs::, 128, 128, 16, 1>::VALID` failed here + +note: erroneous constant encountered + --> src/sm100_attention.rs + | + | let () = Self::VALID; + | ^^^^^^^^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs new file mode 100644 index 0000000000..7012747c86 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.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::TmaDesc; + +fn prefetch(descriptor: &TmaDesc) { + descriptor.prefetch(); +} + +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.stderr new file mode 100644 index 0000000000..b0350be753 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.stderr @@ -0,0 +1,7 @@ +error[E0133]: call to unsafe function `TmaDesc::::prefetch` is unsafe and requires unsafe block + --> tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs:9:5 + | +9 | descriptor.prefetch(); + | ^^^^^^^^^^^^^^^^^^^^^ call to unsafe function + | + = note: consult the function's documentation for information on how to avoid undefined behavior diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.rs new file mode 100644 index 0000000000..888517576d --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.rs @@ -0,0 +1,18 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use cute_rs::sm100_attention::pipeline::{ + AsyncCompletion, Consumer, Producer, RingCursor, TwoCtaRing, +}; + +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()); + } +} +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.stderr new file mode 100644 index 0000000000..d7b48db9e5 --- /dev/null +++ b/cuda-oxide/cuteir/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:12:29 + | +12 | 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, cute_rs::sm100_attention::pipeline::Producer>` + found struct `RingCursor<15, cute_rs::sm100_attention::pipeline::Consumer>` +note: method defined here + --> src/sm100_attention/pipeline.rs + | + | pub unsafe fn acquire_cursor(self, cursor: RingCursor) { + | ^^^^^^^^^^^^^^ + +error[E0308]: mismatched types + --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:13:26 + | +13 | 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, cute_rs::sm100_attention::pipeline::Consumer>` + found struct `RingCursor<15, cute_rs::sm100_attention::pipeline::Producer>` +note: method defined here + --> src/sm100_attention/pipeline.rs + | + | pub unsafe fn wait_cursor(self, cursor: RingCursor) { + | ^^^^^^^^^^^ + +error[E0308]: mismatched types + --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:14:30 + | +14 | 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, cute_rs::sm100_attention::pipeline::Producer>` + found struct `RingCursor<15, cute_rs::sm100_attention::pipeline::Consumer>` +note: method defined here + --> src/sm100_attention/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:15:31 + | +15 | 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, cute_rs::sm100_attention::pipeline::Consumer>` + found struct `RingCursor<15, cute_rs::sm100_attention::pipeline::Producer>` +note: method defined here + --> src/sm100_attention/pipeline.rs + | + | pub unsafe fn mma_empty_cursor(self, cursor: RingCursor) + | ^^^^^^^^^^^^^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs new file mode 100644 index 0000000000..f7cdb0fa3a --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.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::sm100_attention::pipeline::{ + AsyncCompletion, Consumer, Producer, RingCursor, TwoCtaRing, +}; + +unsafe fn wrong_stages(ring: TwoCtaRing<15, AsyncCompletion>) { + unsafe { + ring.acquire_cursor(RingCursor::<2, Producer>::new()); + ring.wait_cursor(RingCursor::<2, Consumer>::new()); + } +} +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.stderr new file mode 100644 index 0000000000..f08ac78904 --- /dev/null +++ b/cuda-oxide/cuteir/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:12:29 + | +12 | 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/sm100_attention/pipeline.rs + | + | pub unsafe fn acquire_cursor(self, cursor: RingCursor) { + | ^^^^^^^^^^^^^^ + +error[E0308]: mismatched types + --> tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs:13:26 + | +13 | 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/sm100_attention/pipeline.rs + | + | pub unsafe fn wait_cursor(self, cursor: RingCursor) { + | ^^^^^^^^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.rs new file mode 100644 index 0000000000..fb3cf1e50a --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.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::sm100_attention::pipeline::{Producer, RingCursor}; + +const _: RingCursor<0, Producer> = RingCursor::new(); +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.stderr new file mode 100644 index 0000000000..aa48eb812d --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.stderr @@ -0,0 +1,12 @@ +error[E0080]: evaluation panicked: assertion failed: STAGES > 0 + --> src/sm100_attention/pipeline.rs + | + | impl RingCursor { + | const VALID: () = assert!(STAGES > 0); + | ^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::sm100_attention::pipeline::RingCursor::<0, cute_rs::sm100_attention::pipeline::Producer>::VALID` failed here + +note: erroneous constant encountered + --> src/sm100_attention/pipeline.rs + | + | let () = Self::VALID; + | ^^^^^^^^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs new file mode 100644 index 0000000000..d92b96f328 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs @@ -0,0 +1,25 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use cute_rs::sm100_attention::pipeline::{ + AsyncCompletion, AsyncCopyReuse, Consumer, LocalShared, RingCursor, TwoCtaRing, +}; + +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.expect(0, 16); + pages.expect(0, 16); + pages.mma_full(0); + pages.mma_empty_cursor(RingCursor::<4, Consumer>::new()); + pages.thread_full(0); + } +} + +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.stderr new file mode 100644 index 0000000000..93e2b2e79c --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.stderr @@ -0,0 +1,149 @@ +error[E0277]: the trait bound `AsyncCopyReuse: SignalPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:16:16 + | +16 | copies.thread_empty(0); + | ^^^^^^^^^^^^ the trait `SignalPolicy` is not implemented for `AsyncCopyReuse` + | +help: the following other types implement trait `SignalPolicy` + --> src/sm100_attention/pipeline.rs + | + | impl SignalPolicy for TensorMemory { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `TensorMemory` +... + | impl SignalPolicy for SharedToAsync { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `SharedToAsync` +note: required by a bound in `TwoCtaRing::::thread_empty` + --> src/sm100_attention/pipeline.rs + | + | impl TwoCtaRing { +... + | pub unsafe fn thread_empty(self, ticket: u32) + | 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:17:15 + | +17 | reuse.expect(0, 16); + | ^^^^^^ the trait `TmaPolicy` is not implemented for `AsyncCopyReuse` + | +help: the trait `TmaPolicy` is implemented for `AsyncCompletion` + --> src/sm100_attention/pipeline.rs + | + | impl TmaPolicy for AsyncCompletion {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +note: required by a bound in `TwoCtaRing::::expect` + --> src/sm100_attention/pipeline.rs + | + | impl TwoCtaRing { +... + | pub unsafe fn expect(self, ticket: u32, bytes: u32) + | where + | Full: TmaPolicy, + | ^^^^^^^^^ required by this bound in `TwoCtaRing::::expect` + +error[E0277]: the trait bound `LocalShared: TmaPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:18:15 + | +18 | pages.expect(0, 16); + | ^^^^^^ the trait `TmaPolicy` is not implemented for `LocalShared` + | +help: the trait `TmaPolicy` is implemented for `AsyncCompletion` + --> src/sm100_attention/pipeline.rs + | + | impl TmaPolicy for AsyncCompletion {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +note: required by a bound in `TwoCtaRing::::expect` + --> src/sm100_attention/pipeline.rs + | + | impl TwoCtaRing { +... + | pub unsafe fn expect(self, ticket: u32, bytes: u32) + | where + | Full: TmaPolicy, + | ^^^^^^^^^ required by this bound in `TwoCtaRing::::expect` + +error[E0277]: the trait bound `LocalShared: MmaPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:19:15 + | +19 | pages.mma_full(0); + | ^^^^^^^^ the trait `MmaPolicy` is not implemented for `LocalShared` + | +help: the following other types implement trait `MmaPolicy` + --> src/sm100_attention/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/sm100_attention/pipeline.rs + | + | impl TwoCtaRing { +... + | pub unsafe fn mma_full(self, ticket: u32) + | 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:20:15 + | +20 | 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/sm100_attention/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/sm100_attention/pipeline.rs + | + | impl TwoCtaRing { +... + | pub unsafe fn mma_empty_cursor(self, cursor: RingCursor) + | 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:21:15 + | +21 | pages.thread_full(0); + | ^^^^^^^^^^^ the trait `SignalPolicy` is not implemented for `LocalShared` + | +help: the following other types implement trait `SignalPolicy` + --> src/sm100_attention/pipeline.rs + | + | impl SignalPolicy for TensorMemory { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `TensorMemory` +... + | impl SignalPolicy for SharedToAsync { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `SharedToAsync` +note: required by a bound in `TwoCtaRing::::thread_full` + --> src/sm100_attention/pipeline.rs + | + | impl TwoCtaRing { +... + | pub unsafe fn thread_full(self, ticket: u32) + | where + | Full: SignalPolicy, + | ^^^^^^^^^^^^ required by this bound in `TwoCtaRing::::thread_full` diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.rs new file mode 100644 index 0000000000..e672adac2d --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.rs @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use cuda_device::barrier::Barrier; +use cute_rs::sm100_attention::pipeline::WaitPolicy; + +#[derive(Clone, Copy)] +struct Unordered; + +impl WaitPolicy for Unordered { + unsafe fn wait(_: *const Barrier, _: u32) {} +} + +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.stderr new file mode 100644 index 0000000000..20b3b6a35b --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.stderr @@ -0,0 +1,33 @@ +error[E0277]: the trait bound `Unordered: cute_rs::sm100_attention::pipeline::sealed::Wait` is not satisfied + --> tests/compile_fail/two_cta_ring_policy_sealed.rs:12:21 + | +12 | impl WaitPolicy for Unordered { + | ^^^^^^^^^ unsatisfied trait bound + | +help: the trait `cute_rs::sm100_attention::pipeline::sealed::Wait` is not implemented for `Unordered` + --> tests/compile_fail/two_cta_ring_policy_sealed.rs:10:1 + | +10 | struct Unordered; + | ^^^^^^^^^^^^^^^^ + = help: the following other types implement trait `cute_rs::sm100_attention::pipeline::sealed::Wait`: + AsyncCompletion + AsyncCopyReuse + LocalMemory + LocalShared + SharedReuse + SharedToAsync + cute_rs::sm100_attention::pipeline::TensorMemory +note: required by a bound in `WaitPolicy` + --> src/sm100_attention/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::sm100_attention::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::sm100_attention::pipeline::LocalMemory + cute_rs::sm100_attention::pipeline::LocalShared + cute_rs::sm100_attention::pipeline::AsyncCompletion + cute_rs::sm100_attention::pipeline::AsyncCopyReuse + cute_rs::sm100_attention::pipeline::TensorMemory + cute_rs::sm100_attention::pipeline::SharedToAsync + cute_rs::sm100_attention::pipeline::SharedReuse diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.rs new file mode 100644 index 0000000000..a728b1a766 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.rs @@ -0,0 +1,25 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use cute_rs::sm100_attention::pipeline::{ + AsyncCompletion, LocalMemory, SharedReuse, SharedToAsync, TensorMemory, TwoCtaRing, +}; + +unsafe fn wrong_completion_mechanisms( + tensor: TwoCtaRing<2, TensorMemory>, + shared: TwoCtaRing<2, SharedToAsync, SharedReuse>, + copies: TwoCtaRing<2, AsyncCompletion>, + local: TwoCtaRing<2, LocalMemory>, +) { + unsafe { + tensor.expect(0, 16); + shared.mma_full(0); + shared.thread_empty(0); + copies.thread_full(0); + local.mma_empty(0); + } +} + +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.stderr new file mode 100644 index 0000000000..ad7ded01d2 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.stderr @@ -0,0 +1,128 @@ +error[E0277]: the trait bound `TensorMemory: TmaPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_wrong_signal.rs:17:16 + | +17 | tensor.expect(0, 16); + | ^^^^^^ the trait `TmaPolicy` is not implemented for `TensorMemory` + | +help: the trait `TmaPolicy` is implemented for `AsyncCompletion` + --> src/sm100_attention/pipeline.rs + | + | impl TmaPolicy for AsyncCompletion {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +note: required by a bound in `TwoCtaRing::::expect` + --> src/sm100_attention/pipeline.rs + | + | impl TwoCtaRing { +... + | pub unsafe fn expect(self, ticket: u32, bytes: u32) + | where + | Full: TmaPolicy, + | ^^^^^^^^^ required by this bound in `TwoCtaRing::::expect` + +error[E0277]: the trait bound `SharedToAsync: MmaPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_wrong_signal.rs:18:16 + | +18 | shared.mma_full(0); + | ^^^^^^^^ the trait `MmaPolicy` is not implemented for `SharedToAsync` + | +help: the following other types implement trait `MmaPolicy` + --> src/sm100_attention/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/sm100_attention/pipeline.rs + | + | impl TwoCtaRing { +... + | pub unsafe fn mma_full(self, ticket: u32) + | 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:19:16 + | +19 | shared.thread_empty(0); + | ^^^^^^^^^^^^ the trait `SignalPolicy` is not implemented for `SharedReuse` + | +help: the following other types implement trait `SignalPolicy` + --> src/sm100_attention/pipeline.rs + | + | impl SignalPolicy for TensorMemory { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `TensorMemory` +... + | impl SignalPolicy for SharedToAsync { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `SharedToAsync` +note: required by a bound in `TwoCtaRing::::thread_empty` + --> src/sm100_attention/pipeline.rs + | + | impl TwoCtaRing { +... + | pub unsafe fn thread_empty(self, ticket: u32) + | 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:20:16 + | +20 | copies.thread_full(0); + | ^^^^^^^^^^^ the trait `SignalPolicy` is not implemented for `AsyncCompletion` + | +help: the following other types implement trait `SignalPolicy` + --> src/sm100_attention/pipeline.rs + | + | impl SignalPolicy for TensorMemory { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `TensorMemory` +... + | impl SignalPolicy for SharedToAsync { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `SharedToAsync` +note: required by a bound in `TwoCtaRing::::thread_full` + --> src/sm100_attention/pipeline.rs + | + | impl TwoCtaRing { +... + | pub unsafe fn thread_full(self, ticket: u32) + | 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:21:15 + | +21 | local.mma_empty(0); + | ^^^^^^^^^ the trait `MmaPolicy` is not implemented for `LocalMemory` + | +help: the following other types implement trait `MmaPolicy` + --> src/sm100_attention/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/sm100_attention/pipeline.rs + | + | impl TwoCtaRing { +... + | pub unsafe fn mma_empty(self, ticket: u32) + | where + | Empty: MmaPolicy, + | ^^^^^^^^^ required by this bound in `TwoCtaRing::::mma_empty` diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.rs new file mode 100644 index 0000000000..0e3a496eec --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.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::sm100_attention::pipeline::{AsyncCompletion, TwoCtaRing}; + +const _: TwoCtaRing<0, AsyncCompletion> = unsafe { TwoCtaRing::from_raw(core::ptr::null_mut(), 0) }; + +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.stderr new file mode 100644 index 0000000000..5ca941d8cd --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.stderr @@ -0,0 +1,12 @@ +error[E0080]: evaluation panicked: assertion failed: STAGES > 0 + --> src/sm100_attention/pipeline.rs + | + | impl TwoCtaRing { + | const VALID: () = assert!(STAGES > 0); + | ^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::sm100_attention::pipeline::TwoCtaRing::<0, cute_rs::sm100_attention::pipeline::AsyncCompletion>::VALID` failed here + +note: erroneous constant encountered + --> src/sm100_attention/pipeline.rs + | + | let () = Self::VALID; + | ^^^^^^^^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.rs new file mode 100644 index 0000000000..ab9ead80e5 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.rs @@ -0,0 +1,6 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +fn main() { + let _ = cute_rs::arch::assume_warp_uniform(0); +} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.stderr new file mode 100644 index 0000000000..cd6bc430fe --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.stderr @@ -0,0 +1,7 @@ +error[E0133]: call to unsafe function `assume_warp_uniform` is unsafe and requires unsafe block + --> tests/compile_fail/warp_uniform_unsafe.rs:5:13 + | +5 | let _ = cute_rs::arch::assume_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/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.rs new file mode 100644 index 0000000000..bb435abb77 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.rs @@ -0,0 +1,6 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +fn main() { + let _ = unsafe { cute_rs::arch::assume_warp_uniform(0u64) }; +} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.stderr new file mode 100644 index 0000000000..929edb2930 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.stderr @@ -0,0 +1,18 @@ +error[E0308]: mismatched types + --> tests/compile_fail/warp_uniform_width.rs:5:57 + | +5 | let _ = unsafe { cute_rs::arch::assume_warp_uniform(0u64) }; + | ---------------------------------- ^^^^ expected `u32`, found `u64` + | | + | arguments to this function are incorrect + | +note: function defined here + --> src/arch.rs + | + | pub unsafe fn assume_warp_uniform(value: u32) -> u32 { + | ^^^^^^^^^^^^^^^^^^^ +help: change the type of the numeric literal from `u64` to `u32` + | +5 - let _ = unsafe { cute_rs::arch::assume_warp_uniform(0u64) }; +5 + let _ = unsafe { cute_rs::arch::assume_warp_uniform(0u32) }; + | diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_pass/tma_descriptor_prefetch.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/tma_descriptor_prefetch.rs new file mode 100644 index 0000000000..40fb4e7cb5 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/tma_descriptor_prefetch.rs @@ -0,0 +1,18 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use cute_rs::TmaDesc; + +// Prefetching metadata preserves the existing element and layout types and +// requires no extra layout trait or conversion at the call site. +unsafe fn prefetch(descriptor: &TmaDesc) { + unsafe { descriptor.prefetch() }; +} + +fn main() { + struct Layout; + let _: unsafe fn(&TmaDesc) = prefetch; + let _: unsafe fn(&TmaDesc) = prefetch; +} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_non_tc_receivers.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_non_tc_receivers.rs new file mode 100644 index 0000000000..971eac0bad --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_non_tc_receivers.rs @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use cute_rs::sm100_attention::pipeline::{ + AsyncCompletion, AsyncCopyReuse, Consumer, LocalShared, Producer, RingCursor, TwoCtaRing, +}; + +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.expect(0, 16); + copies.wait(0); + copies.mma_empty(0); + copies.acquire_cursor(producer); + copies.expect_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); + } +} + +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_cursor.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_cursor.rs new file mode 100644 index 0000000000..fc6a4dd959 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_cursor.rs @@ -0,0 +1,32 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use cute_rs::sm100_attention::pipeline::{ + AsyncCompletion, Consumer, Producer, RingCursor, SharedReuse, SharedToAsync, TensorMemory, + TwoCtaRing, +}; + +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.expect_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); + } +} +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_signal.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_signal.rs new file mode 100644 index 0000000000..d80f9114ff --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_signal.rs @@ -0,0 +1,28 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use cute_rs::sm100_attention::pipeline::{ + AsyncCompletion, SharedReuse, SharedToAsync, TensorMemory, TwoCtaRing, +}; + +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.expect(0, 16); + copies.mma_full(0); + copies.mma_empty(0); + } +} + +fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/global_store.rs b/cuda-oxide/cuteir/cute-rs/tests/global_store.rs new file mode 100644 index 0000000000..5dd4bbb5bc --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/global_store.rs @@ -0,0 +1,28 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#[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::; +} + +#[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); +} + +#[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/cuteir/cute-rs/tests/sm100_attention.rs b/cuda-oxide/cuteir/cute-rs/tests/sm100_attention.rs new file mode 100644 index 0000000000..7f72d6f067 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/sm100_attention.rs @@ -0,0 +1,131 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use cute_rs::sm100_attention::{MmaSs, MmaTs, SharedLayout, SharedTileArray, SmemLayout}; +use cute_rs::{OperandA, OperandB}; + +fn check_bijection() { + let mut offsets = Vec::new(); + for row in 0..L::ROWS { + for col in 0..L::COLS { + offsets.push(L::offset(row, col)); + } + } + offsets.sort_unstable(); + assert_eq!( + offsets, + (0..(L::ROWS * L::COLS) as usize).collect::>() + ); +} + +#[test] +fn tiled_swizzles_cover_storage_without_overlap() { + check_bijection::>(); + check_bijection::>(); + check_bijection::>(); + check_bijection::>(); + check_bijection::>(); +} + +#[test] +fn tma_contiguous_chunks_and_transposed_v_share_the_same_storage() { + type K = SharedLayout<128, 128>; + assert_eq!(K::offset(0, 0), 0); + assert_eq!(K::offset(0, 64), 8192); + assert_eq!(K::offset(1, 0), 72); + for row in 0..32 { + for col in 0..128 { + assert_eq!( + SharedLayout::<32, 128>::offset(row, col), + SharedLayout::<128, 32, true>::offset(col, row) + ); + } + } + assert_eq!(SharedLayout::<64, 32>::SWIZZLE_BITS, 2); + assert_eq!(SharedLayout::<128, 128>::SWIZZLE_BITS, 3); +} + +#[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. +#[repr(align(1024))] +struct SharedStorage([u16; 3 * 16384]); + +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::offset(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::offset(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)); +} + +#[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>(); +} + +#[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. +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/cuteir/cute-rs/tests/sm100_attention_pipeline.rs b/cuda-oxide/cuteir/cute-rs/tests/sm100_attention_pipeline.rs new file mode 100644 index 0000000000..63f2c62a27 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/tests/sm100_attention_pipeline.rs @@ -0,0 +1,87 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use core::mem::MaybeUninit; +use cuda_device::barrier::Barrier; +use cute_rs::sm100_attention::pipeline::{ + AsyncCompletion, Consumer, Producer, RingCursor, TwoCtaRing, +}; + +#[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)); + } + } +} + +#[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"); +} + +#[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)); +} + +#[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(); + } +} + +#[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/cuteir/cute-rs/tests/type_safety.rs b/cuda-oxide/cuteir/cute-rs/tests/type_safety.rs index aeb1b1c552..eba6a93fb2 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/type_safety.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/type_safety.rs @@ -6,6 +6,8 @@ #[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"); @@ -18,4 +20,8 @@ fn tensor_and_block_scaled_types_fail_closed() { 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/cuteir/dialect-cute/src/copy_ops.rs b/cuda-oxide/cuteir/dialect-cute/src/copy_ops.rs new file mode 100644 index 0000000000..730f79965f --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/copy_ops.rs @@ -0,0 +1,130 @@ +/* + * 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}, + location::Located, + 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 { + 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 { + 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 + } + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + 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/cuteir/dialect-cute/src/lib.rs b/cuda-oxide/cuteir/dialect-cute/src/lib.rs index edf6ea76c8..c0970fbac7 100644 --- a/cuda-oxide/cuteir/dialect-cute/src/lib.rs +++ b/cuda-oxide/cuteir/dialect-cute/src/lib.rs @@ -20,12 +20,14 @@ //! The selected backend lowers these operations to GPU code. pub mod attributes; +pub mod copy_ops; pub mod epilogue_ops; pub mod gemm_tma_ops; pub mod gemv_ops; pub mod ops; pub mod pipeline_ops; pub mod scheduler_ops; +pub mod sm100_attention_ops; pub mod sm100_ops; pub mod smem_mma_ops; pub mod tensor_ops; @@ -49,7 +51,10 @@ pub fn register(ctx: &mut Context) { &DialectName::try_new(CUTE_DIALECT_NAME).expect("valid dialect name"), ); + copy_ops::CuteStoreGlobalNoAllocateOp::register(ctx); + copy_ops::CuteGlobalStoreAttr::register(ctx); sm100_ops::register(ctx); + sm100_attention_ops::register(ctx); // The #[pliron_op]/#[pliron_attr] macros auto-register via Context::default; // registering them again is safe and matches the other dialects. diff --git a/cuda-oxide/cuteir/dialect-cute/src/sm100_attention_ops.rs b/cuda-oxide/cuteir/dialect-cute/src/sm100_attention_ops.rs new file mode 100644 index 0000000000..588e85083f --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/sm100_attention_ops.rs @@ -0,0 +1,308 @@ +/* + * 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. +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}, + location::Located, + 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. +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 inner: Layout = format!("{shape}:{stride}") + .parse() + .map_err(|e| format!("{e}"))?; + ComposedLayout::new( + Swizzle::new(if span == 32 { 2 } else { 3 }, 3, 3), + 0, + inner, + OffsetUnit::Elements, + ) + .map(CuteComposedLayoutAttr) + .map_err(|e| e.to_string()) +} + +#[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 { + pub m: u32, + pub n: u32, + pub k: u32, + pub cta_group: u32, + pub a_mn: bool, + pub b_mn: bool, + pub a_layout: CuteComposedLayoutAttr, + 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(()) + } +} +#[pliron_attr(name = "cute.sm100_tma", format = "`<` $layout `,` $cta_group `>`")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteSm100TmaAttr { + pub layout: CuteComposedLayoutAttr, + 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.outer(); + if !matches!(sw.bits, 2 | 3) || sw.base != 3 || sw.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.inner().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(()) + } +} +#[derive(Clone, Copy)] +enum Carrier { + U32, + I32, + Bool, + Shared, + SharedWrite, + Descriptor, + Barrier, +} +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() + } + } +} +macro_rules! operation{ + ($ty:ident,$name:literal,$key:ident,$attr:ty,[$($arg:ident),*])=>{ + #[pliron_op(name=$name,format,attributes=($key:$attr))]pub struct $ty; + impl $ty{ + 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} + pub fn wrap(op:Ptr)->Self{Self{op}} + 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!( + CuteSm100MmaSsOp, + "cute.sm100_mma_ss", + sm100_attention_mma_ss_plan, + CuteSm100MmaAttr, + [U32, Shared, Shared, Bool, U32, U32] +); +operation!( + CuteSm100MmaTsOp, + "cute.sm100_mma_ts", + sm100_attention_mma_ts_plan, + CuteSm100MmaAttr, + [U32, U32, Shared, Bool, U32] +); +operation!( + CuteSm100TmaLoadOp, + "cute.sm100_attention_tma_load", + sm100_attention_tma_load_plan, + CuteSm100TmaAttr, + [Descriptor, SharedWrite, I32, I32, I32, Barrier, U32] +); +operation!( + CuteSm100TmaStoreOp, + "cute.sm100_attention_tma_store", + sm100_attention_tma_store_plan, + CuteSm100TmaAttr, + [Descriptor, Shared, I32, I32, I32] +); +macro_rules! each_op { + ($m:ident) => { + $m!( + CuteSm100MmaSsOp, + CuteSm100MmaTsOp, + CuteSm100TmaLoadOp, + CuteSm100TmaStoreOp + ) + }; +} +pub fn register(ctx: &mut Context) { + macro_rules! reg{($($t:ty),*)=>{$(<$t>::register(ctx);)*}} + each_op!(reg); + CuteSm100MmaAttr::register(ctx); + CuteSm100TmaAttr::register(ctx); +} +pub fn semantic_ids() -> Vec { + macro_rules! ids{($($t:ty),*)=>{vec![$(<$t>::get_opid_static()),*]}} + each_op!(ids) +} +pub fn verify_local(ctx: &Context, op: Ptr) -> Option> { + let id = Operation::get_opid(op, ctx); + 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/cuteir/dialect-cute/src/verify.rs b/cuda-oxide/cuteir/dialect-cute/src/verify.rs index 577231e8e9..2312fd6fef 100644 --- a/cuda-oxide/cuteir/dialect-cute/src/verify.rs +++ b/cuda-oxide/cuteir/dialect-cute/src/verify.rs @@ -199,6 +199,7 @@ fn epilogue_semantic_ids() -> [OpId; 9] { fn all_high_level_semantic_ids() -> Vec { let mut ids = tensor_semantic_ids().to_vec(); + ids.push(crate::copy_ops::CuteStoreGlobalNoAllocateOp::get_opid_static()); ids.extend(gemv_semantic_ids()); ids.extend(scheduler_semantic_ids()); ids.extend(pipeline_semantic_ids()); @@ -206,6 +207,7 @@ fn all_high_level_semantic_ids() -> Vec { ids.extend(smem_mma_semantic_ids()); ids.extend(epilogue_semantic_ids()); ids.extend(crate::sm100_ops::semantic_ids()); + ids.extend(crate::sm100_attention_ops::semantic_ids()); ids } @@ -217,7 +219,11 @@ fn verify_local_op(ctx: &Context, operation: Ptr) -> Result { }; } - let result = if let Some(result) = crate::sm100_ops::verify_local(ctx, operation) { + let result = if opid == crate::copy_ops::CuteStoreGlobalNoAllocateOp::get_opid_static() { + Some(verify_as!(crate::copy_ops::CuteStoreGlobalNoAllocateOp)) + } else if let Some(result) = crate::sm100_attention_ops::verify_local(ctx, operation) { + Some(result) + } else if let Some(result) = crate::sm100_ops::verify_local(ctx, operation) { Some(result) } else if opid == CuteTensorMakeOp::get_opid_static() { Some(verify_as!(CuteTensorMakeOp)) @@ -3838,6 +3844,7 @@ pub fn verify_cute_semantics(ctx: &Context, module: Ptr) -> Result { verify_tensor_story(ctx, &scope)?; verify_epilogue_story(ctx, &scope)?; crate::sm100_ops::verify_story(ctx, &scope).map_err(invalid)?; + crate::sm100_attention_ops::verify_story(ctx, &scope).map_err(invalid)?; } Ok(()) } diff --git a/cuda-oxide/cuteir/dialect-cute/tests/global_store.rs b/cuda-oxide/cuteir/dialect-cute/tests/global_store.rs new file mode 100644 index 0000000000..24548e8e8d --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/tests/global_store.rs @@ -0,0 +1,54 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use dialect_cute::copy_ops::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, +}; + +#[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/cuteir/dialect-cute/tests/sm100_attention.rs b/cuda-oxide/cuteir/dialect-cute/tests/sm100_attention.rs new file mode 100644 index 0000000000..66568fc848 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/tests/sm100_attention.rs @@ -0,0 +1,165 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use dialect_cute::sm100_attention_ops::{CuteSm100MmaAttr, shared_layout}; +use pliron::{common_traits::Verify, context::Context}; +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(), + } +} +#[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()); +} +#[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()); +} +#[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()); +} + +#[test] +fn tma_rejects_noncanonical_layout() { + use dialect_cute::attributes::CuteComposedLayoutAttr; + use dialect_cute::layout::{ComposedLayout, OffsetUnit, Swizzle}; + use dialect_cute::sm100_attention_ops::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() + ); +} + +#[test] +fn tma_load_rejects_readonly_shared_destination() { + use dialect_cute::sm100_attention_ops::{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); + } +} + +#[test] +fn mma_tile_offsets_require_explicit_u32_operands() { + use dialect_cute::sm100_attention_ops::{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/cuteir/docs/decode-parity-on-sm100.md b/cuda-oxide/cuteir/docs/decode-parity-on-sm100.md new file mode 100644 index 0000000000..b39bbfe853 --- /dev/null +++ b/cuda-oxide/cuteir/docs/decode-parity-on-sm100.md @@ -0,0 +1,5 @@ + +In a nutshell + +- **Smarter synchronization and bookkeeping:** Added compiler support and typed pipelines that avoid unnecessary cache invalidations. Reused memory descriptors and reduced repeated calculations. +- **Wider output stores:** Added typed, aligned 256-bit copies and their compiler lowering—twice the data per store. This closed the remaining gap while keeping alignment safety in Rust’s types. \ No newline at end of file diff --git a/cuda-oxide/cuteir/docs/translation.md b/cuda-oxide/cuteir/docs/translation.md index 44f16671cb..a0093e0001 100644 --- a/cuda-oxide/cuteir/docs/translation.md +++ b/cuda-oxide/cuteir/docs/translation.md @@ -33,6 +33,19 @@ CUDA_OXIDE_CUTLASS_COMPILER=/absolute/path/to/libCutlassCompiler.so \ cargo oxide build elementwise_cute --arch sm_120a ``` +Attention comparisons use the public reference's optimization level 2: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 \ + cargo oxide build sm100_fmha_prefill_cute --arch sm_100a +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 \ + cargo oxide build sm100_mla_decode_fp16_cute --arch sm_100a +``` + +`CUDA_OXIDE_CUTLASS_OPT_LEVEL` accepts 2 or 3, with 3 remaining the default. +The execution manifest records the chosen level and the resulting pipeline +hash; changing the environment option invalidates the device build fingerprint. + ## Supported mappings | Example | Compiler-visible operations | @@ -42,6 +55,15 @@ CUDA_OXIDE_CUTLASS_COMPILER=/absolute/path/to/libCutlassCompiler.so \ | [Block-scaled GEMM](../examples/blockscale_gemm_cute) | Scheduler, work tiles, TMA pipelines, shared-memory MMA, epilogue stores | | [SM100 FP16 GEMM](../examples/fp16_gemm_256x352_cute) | Typed shared tiles, TMEM, paired two-CTA MMA, cluster TMA, producer/consumer pipelines | +The [SM100 attention API](../cute-rs/src/sm100_attention.rs) supplies independent +QK and PV plans: group-one/group-two shared/shared MMA, group-one TMEM/shared MMA, +K-major and MN-major tiled shared layouts, rank-three typed TMA copies, and +FP32 TMEM fragments with packed FP16 stores. The semantic verifier checks +operand roles, layout extents, major order, and MMA group compatibility. +Attention kernels keep their warp schedule, online softmax, barrier phases, +and asynchronous buffer lifetimes explicit. Those dynamic obligations remain +with unsafe callers; a local operation order does not prove a cross-warp handoff. + At a block-scaled epilogue hand-off, `ReadyForTma` makes shared-memory writes visible to TMA with an async-shared proxy fence, then emits a counted CTA barrier. `Reusable` emits only the barrier. diff --git a/cuda-oxide/cuteir/examples/blockscale_gemm_cute/Cargo.lock b/cuda-oxide/cuteir/examples/blockscale_gemm_cute/Cargo.lock index b772601478..219d535376 100644 --- a/cuda-oxide/cuteir/examples/blockscale_gemm_cute/Cargo.lock +++ b/cuda-oxide/cuteir/examples/blockscale_gemm_cute/Cargo.lock @@ -125,6 +125,13 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" @@ -213,6 +220,7 @@ version = "0.1.0" name = "cute-rs" version = "0.1.0" dependencies = [ + "ctm-rs", "cuda-core", "cuda-device", "cute-layout", diff --git a/cuda-oxide/cuteir/examples/elementwise_cute/Cargo.lock b/cuda-oxide/cuteir/examples/elementwise_cute/Cargo.lock index 736c2ac36f..3500eca157 100644 --- a/cuda-oxide/cuteir/examples/elementwise_cute/Cargo.lock +++ b/cuda-oxide/cuteir/examples/elementwise_cute/Cargo.lock @@ -115,6 +115,13 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" @@ -203,6 +210,7 @@ version = "0.1.0" name = "cute-rs" version = "0.1.0" dependencies = [ + "ctm-rs", "cuda-device", "cute-layout", ] diff --git a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/Cargo.lock b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/Cargo.lock index c788c26f9d..661f5b8b98 100644 --- a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/Cargo.lock +++ b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/Cargo.lock @@ -115,6 +115,13 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" @@ -203,6 +210,7 @@ version = "0.1.0" name = "cute-rs" version = "0.1.0" dependencies = [ + "ctm-rs", "cuda-device", "cute-layout", ] diff --git a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/Cargo.lock b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/Cargo.lock index 126668a3b2..f76b544bbb 100644 --- a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/Cargo.lock +++ b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/Cargo.lock @@ -115,6 +115,13 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" @@ -203,6 +210,7 @@ version = "0.1.0" name = "cute-rs" version = "0.1.0" dependencies = [ + "ctm-rs", "cuda-device", "cute-layout", ] diff --git a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/Cargo.lock b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/Cargo.lock index fb774d4e76..94cbbaf89f 100644 --- a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/Cargo.lock +++ b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/Cargo.lock @@ -210,6 +210,7 @@ version = "0.1.0" name = "cute-rs" version = "0.1.0" dependencies = [ + "ctm-rs", "cuda-core", "cuda-device", "cute-layout", diff --git a/cuda-oxide/cuteir/examples/nvfp4_gemv_cute/Cargo.lock b/cuda-oxide/cuteir/examples/nvfp4_gemv_cute/Cargo.lock index 67f4dfbc52..64698b3380 100644 --- a/cuda-oxide/cuteir/examples/nvfp4_gemv_cute/Cargo.lock +++ b/cuda-oxide/cuteir/examples/nvfp4_gemv_cute/Cargo.lock @@ -115,6 +115,13 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" @@ -203,6 +210,7 @@ version = "0.1.0" name = "cute-rs" version = "0.1.0" dependencies = [ + "ctm-rs", "cuda-device", "cute-layout", ] diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/Cargo.lock b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/Cargo.lock new file mode 100644 index 0000000000..ea6668c0e8 --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/Cargo.lock @@ -0,0 +1,779 @@ +# 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.4.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-core" +version = "0.4.0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.4.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 = [ + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn 2.0.119", +] + +[[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.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" + +[[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.190" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce5d3ddc6d3fa000eb1536d85e147bfe31aacaba692ed6a876f95cb7c855be78" + +[[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 = "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 = "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.8", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+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.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.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01016da373cd8f7ef12624f796309f5c31ba8d646dd08856c02cd741d823c622" +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.8", +] + +[[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.62" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "86502bf56ac7c77571a32e2647bb2a15894565e981fb2a48d7bde2d91c965a9d" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.62" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5457206954b06561e2608c7e19cf58b1926586d999c246eebe4502f7e2039d1a" +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/cuteir/examples/sm100_fmha_prefill_cute/Cargo.toml b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/Cargo.toml new file mode 100644 index 0000000000..e7da546e0b --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/examples/sm100_fmha_prefill_cute/README.md b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/README.md new file mode 100644 index 0000000000..2f2154c4c8 --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/README.md @@ -0,0 +1,55 @@ +# 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. + +- 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/`, with CUTLASS 4.7 installed (`cargo oxide toolchain install cutlass`) +and the [Python reference environment](reference/README.md) available: + +```bash +# 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:
];` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) | | `cuda_intrinsics::barrier::mbarrier_arrive_cluster` | `nvvm.mbarrier_arrive_cluster` | PTX-native `mbarrier.arrive.release.cluster.shared::cluster.b64` | unsafe; scope cluster; impure; memory read_write; convergent true | PTX 8.0 on sm_90+ ([PTX ISA 9.3, Parallel Synchronization and Communication Instructions: mbarrier.arrive](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-arrive)) | validated on `rust-llvm-23.1.3-aab11ee0`; expects `mbarrier.arrive.release.cluster.shared::cluster.b64 _, [
];` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) | +| `cuda_intrinsics::barrier::mbarrier_arrive_cluster_relaxed` | `nvvm.mbarrier_arrive_cluster_relaxed` | PTX-native `mbarrier.arrive.relaxed.cluster.shared::cluster.b64` | unsafe; scope cluster; impure; memory read_write; convergent true | PTX 8.6 on sm_90+ ([PTX ISA 9.3, Parallel Synchronization and Communication Instructions: mbarrier.arrive](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-arrive)) | validated on `rust-llvm-23.1.3-aab11ee0`; expects `mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [
];` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) | | `cuda_intrinsics::barrier::mbarrier_arrive_expect_tx` | `nvvm.mbarrier_arrive_expect_tx_shared` | LLVM `llvm.nvvm.mbarrier.arrive.expect.tx.scope.cta.space.cta` from `int_nvvm_mbarrier_arrive_expect_tx_scope_cta_space_cta` | unsafe; scope cta; impure; memory read_write; convergent true | PTX 8.0 on sm_90+ ([PTX ISA 9.3, Parallel Synchronization and Communication Instructions: mbarrier.arrive](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-arrive)) | validated on `rust-llvm-23.1.3-aab11ee0`; expects `mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 , [
], ;` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) | | `cuda_intrinsics::barrier::mbarrier_arrive_expect_tx_cluster` | `nvvm.mbarrier_arrive_expect_tx_cluster` | LLVM `llvm.nvvm.mbarrier.arrive.expect.tx.relaxed.scope.cluster.space.cta` from `int_nvvm_mbarrier_arrive_expect_tx_relaxed_scope_cluster_space_cta` | unsafe; scope cluster; impure; memory read_write; convergent true | PTX 8.6 on sm_90+ ([PTX ISA 9.3, Parallel Synchronization and Communication Instructions: mbarrier.arrive](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-arrive)) | validated on `rust-llvm-23.1.3-aab11ee0`; expects `mbarrier.arrive.expect_tx.relaxed.cluster.shared::cta.b64 , [
], ;` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) | | `cuda_intrinsics::barrier::mbarrier_arrive_no_complete` | `nvvm.mbarrier_arrive_no_complete_shared` | LLVM `llvm.nvvm.mbarrier.arrive.noComplete.shared` from `int_nvvm_mbarrier_arrive_noComplete_shared` | unsafe; scope cta; impure; memory read_write; convergent true | PTX 7.0 on sm_80+ ([PTX ISA 9.3, 9.7.14.16.16 Parallel Synchronization and Communication Instructions: mbarrier.arrive](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-arrive)) | validated on `rust-llvm-23.1.3-aab11ee0`; expects `mbarrier.arrive.noComplete.shared.b64 , [
], ;` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) | @@ -364,6 +366,7 @@ This table is generated from the resolved catalog. Evidence stages below disting | `cuda_intrinsics::barrier::mbarrier_try_wait` | `nvvm.mbarrier_try_wait_shared` | LLVM `llvm.nvvm.mbarrier.try.wait.scope.cta.space.cta` from `int_nvvm_mbarrier_try_wait_scope_cta_space_cta` | unsafe; scope cta; impure; memory read_write; convergent true | PTX 7.8 on sm_90+ ([PTX ISA 9.3, Parallel Synchronization and Communication Instructions: mbarrier.test_wait / mbarrier.try_wait](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-test-wait-mbarrier-try-wait)) | validated on `rust-llvm-23.1.3-aab11ee0`; expects `mbarrier.try_wait.shared.b64 , [
], ;` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) | | `cuda_intrinsics::barrier::mbarrier_try_wait_parity` | `nvvm.mbarrier_try_wait_parity_shared` | LLVM `llvm.nvvm.mbarrier.try.wait.parity.scope.cta.space.cta` from `int_nvvm_mbarrier_try_wait_parity_scope_cta_space_cta` | unsafe; scope cta; impure; memory read_write; convergent true | PTX 7.8 on sm_90+ ([PTX ISA 9.3, Parallel Synchronization and Communication Instructions: mbarrier.test_wait / mbarrier.try_wait](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-test-wait-mbarrier-try-wait)) | validated on `rust-llvm-23.1.3-aab11ee0`; expects `mbarrier.try_wait.parity.shared::cta.b64 , [
], ;` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) | | `cuda_intrinsics::barrier::mbarrier_try_wait_parity_cluster` | `nvvm.mbarrier_try_wait_parity_cluster` | LLVM `llvm.nvvm.mbarrier.try.wait.parity.scope.cluster.space.cta` from `int_nvvm_mbarrier_try_wait_parity_scope_cluster_space_cta` | unsafe; scope cluster; impure; memory read_write; convergent true | PTX 8.0 on sm_90+ ([PTX ISA 9.3, Parallel Synchronization and Communication Instructions: mbarrier.test_wait / mbarrier.try_wait](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-test-wait-mbarrier-try-wait)) | validated on `rust-llvm-23.1.3-aab11ee0`; expects `mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 , [
], ;` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) | +| `cuda_intrinsics::barrier::mbarrier_try_wait_parity_cluster_relaxed` | `nvvm.mbarrier_try_wait_parity_cluster_relaxed` | LLVM `llvm.nvvm.mbarrier.try.wait.parity.relaxed.scope.cluster.space.cta` from `int_nvvm_mbarrier_try_wait_parity_relaxed_scope_cluster_space_cta` | unsafe; scope cluster; impure; memory read_write; convergent true | PTX 8.6 on sm_90+ ([PTX ISA 9.3, Parallel Synchronization and Communication Instructions: mbarrier.test_wait / mbarrier.try_wait](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-test-wait-mbarrier-try-wait)) | validated on `rust-llvm-23.1.3-aab11ee0`; expects `mbarrier.try_wait.parity.relaxed.cluster.shared::cta.b64 , [
], ;` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) | | `cuda_intrinsics::bf16::min_bf16` | `nvvm.extended_minmax` | LLVM `llvm.nvvm.fmin.bf16` from `int_nvvm_fmin_bf16` | safe; scope thread; pure; memory none; convergent false | PTX 7.0 on sm_80+ ([PTX ISA 9.3, 9.7.4.7 Half Precision Floating Point Instructions: min](https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-min)) | validated on `rust-llvm-23.1.3-aab11ee0`; expects `min.bf16 , , ;` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) | | `cuda_intrinsics::bf16x2::min_bf16x2` | `nvvm.min_bf16x2` | LLVM `llvm.nvvm.fmin.bf16x2` from `int_nvvm_fmin_bf16x2` | safe; scope thread; pure; memory none; convergent false | PTX 7.0 on sm_80+ ([PTX ISA 9.3, 9.7.4.7 Half Precision Floating Point Instructions: min](https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-min)) | validated on `rust-llvm-23.1.3-aab11ee0`; expects `min.bf16x2 , , ;` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) | | `cuda_intrinsics::f16::min_f16` | `nvvm.extended_minmax` | LLVM `llvm.nvvm.fmin.f16` from `int_nvvm_fmin_f16` | safe; scope thread; pure; memory none; convergent false | PTX 7.0 on sm_80+ ([PTX ISA 9.3, 9.7.4.7 Half Precision Floating Point Instructions: min](https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-min)) | validated on `rust-llvm-23.1.3-aab11ee0`; expects `min.f16 , , ;` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) | @@ -1965,6 +1968,10 @@ This table is generated from the resolved catalog. Evidence stages below disting - canonical compiler path: `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0820` - public source path: `cuda_intrinsics::float::ex2_approx_ftz_f32` - compatibility compiler path: `cuda_device::float::ex2_approx_ftz_f32` +- `fence_acquire_shared_cluster_cluster` (`i1033`) + - canonical compiler path: `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033` + - public source path: `cuda_intrinsics::barrier::fence_acquire_shared_cluster_cluster` + - compatibility compiler path: `cuda_device::barrier::fence_acquire_shared_cluster_cluster` - `fence_mbarrier_init_release_cluster` (`i0313`) - canonical compiler path: `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0313` - public source path: `cuda_intrinsics::barrier::fence_mbarrier_init_release_cluster` @@ -2422,6 +2429,10 @@ This table is generated from the resolved catalog. Evidence stages below disting - canonical compiler path: `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0308` - public source path: `cuda_intrinsics::barrier::mbarrier_arrive_cluster` - compatibility compiler path: `cuda_device::barrier::mbarrier_arrive_cluster` +- `mbarrier_arrive_cluster_relaxed` (`i1031`) + - canonical compiler path: `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1031` + - public source path: `cuda_intrinsics::barrier::mbarrier_arrive_cluster_relaxed` + - compatibility compiler path: `cuda_device::barrier::mbarrier_arrive_cluster_relaxed` - `mbarrier_arrive_expect_tx` (`i0306`) - canonical compiler path: `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0306` - public source path: `cuda_intrinsics::barrier::mbarrier_arrive_expect_tx` @@ -2458,6 +2469,10 @@ This table is generated from the resolved catalog. Evidence stages below disting - canonical compiler path: `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0311` - public source path: `cuda_intrinsics::barrier::mbarrier_try_wait_parity_cluster` - compatibility compiler path: `cuda_device::barrier::mbarrier_try_wait_parity_cluster` +- `mbarrier_try_wait_parity_cluster_relaxed` (`i1032`) + - canonical compiler path: `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1032` + - public source path: `cuda_intrinsics::barrier::mbarrier_try_wait_parity_cluster_relaxed` + - compatibility compiler path: `cuda_device::barrier::mbarrier_try_wait_parity_cluster_relaxed` - `min_bf16` (`i0838`) - canonical compiler path: `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0838` - public source path: `cuda_intrinsics::bf16::min_bf16` @@ -5682,16 +5697,19 @@ This table is generated from the resolved catalog. Evidence stages below disting ## Extended mbarrier contracts +- `fence_acquire_shared_cluster_cluster` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is LLVM 22 TableGen. - `fence_mbarrier_init_release_cluster` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is LLVM 22 TableGen. - `fence_proxy_async_generic_acquire_shared_cluster_cluster` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is LLVM 22 TableGen. - `fence_proxy_async_generic_release_shared_cta_cluster` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is LLVM 22 TableGen. - `fence_proxy_async_shared_cta` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is LLVM 22 TableGen. - `mbarrier_arrive_cluster` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is PTX-native raw-address contract. +- `mbarrier_arrive_cluster_relaxed` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is PTX-native raw-address contract. - `mbarrier_arrive_expect_tx` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is LLVM 22 TableGen. - `mbarrier_arrive_expect_tx_cluster` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is LLVM 22 TableGen. - `mbarrier_try_wait` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is LLVM 22 TableGen. - `mbarrier_try_wait_parity` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is LLVM 22 TableGen. - `mbarrier_try_wait_parity_cluster` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is LLVM 22 TableGen. +- `mbarrier_try_wait_parity_cluster_relaxed` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is LLVM 22 TableGen. - `nanosleep` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is LLVM 22 TableGen. ## WGMMA-control contracts @@ -5953,6 +5971,7 @@ This table is generated from the resolved catalog. Evidence stages below disting - `ex2_approx_f16`: properties `IntrNoMem`; result `NoUndef` false; half-open result range `none`; selection records ; selection predicates none - `ex2_approx_f32`: properties `IntrNoMem`; result `NoUndef` false; half-open result range `none`; selection records ; selection predicates none - `ex2_approx_ftz_f32`: properties `IntrNoMem`; result `NoUndef` false; half-open result range `none`; selection records ; selection predicates none +- `fence_acquire_shared_cluster_cluster`: properties `IntrNoCallback`; result `NoUndef` false; half-open result range `none`; selection records `INT_FENCE_ACQUIRE_SYNC_RESTRICT_CLUSTER_CLUSTER`; selection predicates `Subtarget->getPTXVersion() >= 86`, `Subtarget->getSmVersion() >= 90` - `fence_mbarrier_init_release_cluster`: properties `IntrNoCallback`; result `NoUndef` false; half-open result range `none`; selection records `INT_FENCE_MBARRIER_INIT_RELEASE_CLUSTER`; selection predicates `Subtarget->getPTXVersion() >= 80`, `Subtarget->getSmVersion() >= 90` - `fence_proxy_async_generic_acquire_shared_cluster_cluster`: properties `IntrNoCallback`; result `NoUndef` false; half-open result range `none`; selection records `INT_NVVM_FENCE_PROXY_ASYNC_GENERIC_ACQUIRE_SYNC_RESTRICT_SPACE_CLUSTER_SCOPE_CLUSTER`; selection predicates `Subtarget->getPTXVersion() >= 86`, `Subtarget->getSmVersion() >= 90` - `fence_proxy_async_generic_release_shared_cta_cluster`: properties `IntrNoCallback`; result `NoUndef` false; half-open result range `none`; selection records `INT_NVVM_FENCE_PROXY_ASYNC_GENERIC_RELEASE_SYNC_RESTRICT_SPACE_CTA_SCOPE_CLUSTER`; selection predicates `Subtarget->getPTXVersion() >= 86`, `Subtarget->getSmVersion() >= 90` @@ -6067,6 +6086,7 @@ This table is generated from the resolved catalog. Evidence stages below disting - `max_xorsign_abs_f32`: properties `Commutative`, `IntrNoCreateUndefOrPoison`, `IntrNoMem`, `IntrSpeculatable`; result `NoUndef` false; half-open result range `none`; selection records `INT_NVVM_FMAX_XORSIGN_ABS_F`; selection predicates `Subtarget->getPTXVersion() >= 72`, `Subtarget->getSmVersion() >= 86` - `mbarrier_arrive`: properties `IntrConvergent`, `IntrNoCallback`; result `NoUndef` false; half-open result range `none`; selection records `MBARRIER_ARRIVE_SHARED`; selection predicates `Subtarget->getPTXVersion() >= 70`, `Subtarget->getSmVersion() >= 80` - `mbarrier_arrive_cluster`: PTX-native source; no LLVM record, symbol, properties, or selection facts. +- `mbarrier_arrive_cluster_relaxed`: PTX-native source; no LLVM record, symbol, properties, or selection facts. - `mbarrier_arrive_expect_tx`: properties `IntrConvergent`, `IntrNoCallback`; result `NoUndef` false; half-open result range `none`; selection records ; selection predicates none - `mbarrier_arrive_expect_tx_cluster`: properties `IntrArgMemOnly`, `IntrConvergent`, `IntrNoCallback`; result `NoUndef` false; half-open result range `none`; selection records ; selection predicates none - `mbarrier_arrive_no_complete`: properties `IntrConvergent`, `IntrNoCallback`; result `NoUndef` false; half-open result range `none`; selection records `MBARRIER_ARRIVE_NOCOMPLETE_SHARED`; selection predicates `Subtarget->getPTXVersion() >= 70`, `Subtarget->getSmVersion() >= 80` @@ -6076,6 +6096,7 @@ This table is generated from the resolved catalog. Evidence stages below disting - `mbarrier_try_wait`: properties `IntrConvergent`, `IntrNoCallback`, `NoCapture`; result `NoUndef` false; half-open result range `none`; selection records `mbar_try_wait_scope_cta_acquire_STATE`; selection predicates `Subtarget->getPTXVersion() >= 78`, `Subtarget->getSmVersion() >= 90` - `mbarrier_try_wait_parity`: properties `IntrConvergent`, `IntrNoCallback`, `NoCapture`; result `NoUndef` false; half-open result range `none`; selection records ; selection predicates none - `mbarrier_try_wait_parity_cluster`: properties `IntrConvergent`, `IntrNoCallback`, `NoCapture`; result `NoUndef` false; half-open result range `none`; selection records ; selection predicates none +- `mbarrier_try_wait_parity_cluster_relaxed`: properties `IntrArgMemOnly`, `IntrConvergent`, `IntrNoCallback`, `IntrReadMem`, `NoCapture`; result `NoUndef` false; half-open result range `none`; selection records ; selection predicates none - `min_bf16`: properties `Commutative`, `IntrNoCreateUndefOrPoison`, `IntrNoMem`, `IntrSpeculatable`; result `NoUndef` false; half-open result range `none`; selection records `INT_NVVM_FMIN_bf16`; selection predicates `Subtarget->getPTXVersion() >= 70`, `Subtarget->getSmVersion() >= 80` - `min_bf16x2`: properties `Commutative`, `IntrNoCreateUndefOrPoison`, `IntrNoMem`, `IntrSpeculatable`; result `NoUndef` false; half-open result range `none`; selection records `INT_NVVM_FMIN_bf16x2`; selection predicates `Subtarget->getPTXVersion() >= 70`, `Subtarget->getSmVersion() >= 80` - `min_f16`: properties `Commutative`, `IntrNoCreateUndefOrPoison`, `IntrNoMem`, `IntrSpeculatable`; result `NoUndef` false; half-open result range `none`; selection records `INT_NVVM_FMIN_f16`; selection predicates `Subtarget->getPTXVersion() >= 70`, `Subtarget->getSmVersion() >= 80` @@ -8519,6 +8540,13 @@ This table is generated from the resolved catalog. Evidence stages below disting - backend codegen on `compute_80, ptx93`: `succeeded` — CUDA 13.3 libNVVM verified and compiled a dedicated typed-pointer NVVM IR fixture holding the 3 appended variants (ex2.approx.f32, ex2.approx.ftz.f32, tanh.approx.f32) for compute_80. The emitted PTX 9.3 output retains each instruction exactly once with no extern calls; PTX SHA-256 9d2a3b28989cf40d76816090ce6051c2347f46d8bbbd6fd7fb5b479977e5d7cc. - PTX assembly on `sm_80, ptx93`: `succeeded` — CUDA 13.3 ptxas assembled the appended-fixture PTX into an sm_80 cubin with SHA-256 134a765e19dd2267cf41934ad4ffa3ec0a6f54c43406d4da3102adfb9334e955. tanh.approx.f32's hardware floor is sm_75 (PTX ISA 7.0); the family admits it at the attested sm_80 evidence floor. Tool `/usr/local/cuda-13.3/bin/ptxas` reports `Cuda compilation tools, release 13.3, V13.3.33; build cuda_13.3.r13.3/compiler.37862127_0` (SHA-256 `7fdd01a4cf50e30746da98989c9272a907f491e6fd7fecfda14642e4375f88fb`). Artifact `cubin`. - device link on `sm_80, ptx93`: `succeeded` — CUDA 13.3 libNVVM compiled the appended fixture with -gen-lto to a 1,848-byte LTOIR (SHA-256 6c6fcf6eb5ddef3cf3ea605c0a866c8cc4db63bae1df7abca521ff507283b56d) and nvJitLink 13.3 linked it into a 1,464-byte sm_80 cubin (SHA-256 3e4a06eab6319f8d75a890f4b2645d35d92125c79eaa2b63dba2e5383552d632) with all three appended instructions retained. Tool `/usr/local/cuda-13.3/targets/x86_64-linux/lib/libnvJitLink.so.13.3.33` reports `nvJitLinkVersion 13.3; CUDA 13.3 V13.3.33` (SHA-256 `3ba1e744347cd68617b862eccfd98b125482e882b7a6319f42abc9a768513db8`). Artifact `cubin`. +- `fence_acquire_shared_cluster_cluster`: runtime `unexecuted` + - `llvm_nvptx` uses `inline_ptx` from profile `rust-llvm-23.1.3-aab11ee0` at PTX 8.6 / sm_90+: status `validated` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) + - backend codegen on `sm_90, ptx86`: `succeeded` — Compiled and checked the retained exact instruction on 2026-10-12. PTX SHA-256 7630a6dc59639dc67bcd19e386d47b226c2b4a4dbc7ad63ffffb3c69b56ebc83. + - PTX assembly on `sm_90, ptx86`: `succeeded` — CPU terminal validation retained the probe on 2026-10-12. Cubin SHA-256 07eaacbe13a41fb20299adb0fd343b686fbe179cac1a4770c2206421f1eeb630. Tool `/usr/local/cuda/bin/ptxas` reports `Cuda compilation tools, release 13.4, V13.4.92` (SHA-256 `6f6a701573ac9f5c2409708b3b700cb9b2a0f11b13c20468153c13b597b2d6d3`). Artifact `cubin`. + - `lib_nvvm` uses `inline_ptx` from profile `cuda-13.4-libnvvm-13.4.92-mbarrier-relaxed-cluster` at PTX 8.6 / sm_90+: status `validated` (`CUDA toolkit 13.4; nvcc/libNVVM release V13.4.92`, SHA-256 `eeef1ca79dc4ded9cc4ceb69a12f6d7148b5ccb5fc9f997a8ef40ff050be7092`) + - backend codegen on `sm_90, ptx94`: `succeeded` — Compiled and checked the retained exact instruction on 2026-10-12. PTX SHA-256 e6cee64c4899ecdfb58ea4e3a87f6e6c917dba0b489fd7d48a1897fd54ba3b31. + - device link on `sm_90, ptx94`: `succeeded` — CPU terminal validation retained the probe on 2026-10-12. Cubin SHA-256 3d6f528ee6d25eb65d606570926cf4d6eb0f595427a1423b08877d8e7113cde7. Tool `/usr/local/cuda/lib64/libnvJitLink.so.13.4.92` reports `V13.4.92` (SHA-256 `3afdb9677477dd584d3712a9bf2978d085729f0ed6e6432027e925d1e312e3c9`). Artifact `cubin`. - `fence_mbarrier_init_release_cluster`: runtime `unexecuted` - `llvm_nvptx` uses `inline_ptx` from profile `rust-llvm-23.1.3-aab11ee0` at PTX 8.0 / sm_90+: status `validated` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) - backend codegen on `sm_90, ptx80`: `succeeded` — Revalidated 5 separate probes with nightly-2026-10-11 (rustc 6f5da406b56c2db6607202b63daee73aa8c61297; LLVM aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019): the exact reviewed PTX instruction shapes passed. The sorted ID/artifact-kind/SHA-256 ledger in refresh-2026-10-11/stages.jsonl has SHA-256 43181bf08ce9b3b668f72d19f5c2ff1fda31995dbff3b3e517bbe3ede29e620c. @@ -9947,6 +9975,13 @@ This table is generated from the resolved catalog. Evidence stages below disting - `lib_nvvm` uses `inline_ptx` from profile `cuda-13.3-libnvvm-13.3.33-mbarrier-extended` at PTX 8.0 / sm_90+: status `validated` (`CUDA toolkit 13.3; nvcc/libNVVM release V13.3.33`, SHA-256 `b6a5db233133fb0ed693f8cb8aeb1d405254ce0ab10364478985292ec98cfcf0`) - backend codegen on `compute_90, ptx93`: `succeeded` — CUDA 13.3 libNVVM verified the modules and retained each exact extended-mbarrier or async-proxy fence instruction. - device link on `sm_90, ptx93`: `succeeded` — CUDA 13.3 nvJitLink linked the libNVVM-produced LTOIR into cubins for sm_90. Tool `/usr/local/cuda-13.3/lib64/libnvJitLink.so.13.3.33` reports `V13.3.33` (SHA-256 `3ba1e744347cd68617b862eccfd98b125482e882b7a6319f42abc9a768513db8`). Artifact `cubin`. +- `mbarrier_arrive_cluster_relaxed`: runtime `unexecuted` + - `llvm_nvptx` uses `inline_ptx` from profile `rust-llvm-23.1.3-aab11ee0` at PTX 8.6 / sm_90+: status `validated` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) + - backend codegen on `sm_90, ptx86`: `succeeded` — Compiled and checked the exact instruction on 2026-10-12. PTX SHA-256 95a3f51783825dfa9a4baf9218584d33c8af170f20863e86da5ea8f542bcbec8. + - PTX assembly on `sm_90, ptx86`: `succeeded` — ptxas -arch=sm_90 -c accepted the generated PTX on 2026-10-12. Cubin SHA-256 59e3616fa4b0e8d6f51b3b333de30dcac74f1ba002744cb6c6c918cae5ae6463. Tool `/usr/local/cuda/bin/ptxas` reports `Cuda compilation tools, release 13.4, V13.4.92` (SHA-256 `6f6a701573ac9f5c2409708b3b700cb9b2a0f11b13c20468153c13b597b2d6d3`). Artifact `cubin`. + - `lib_nvvm` uses `inline_ptx` from profile `cuda-13.4-libnvvm-13.4.92-mbarrier-relaxed-cluster` at PTX 8.6 / sm_90+: status `validated` (`CUDA toolkit 13.4; nvcc/libNVVM release V13.4.92`, SHA-256 `eeef1ca79dc4ded9cc4ceb69a12f6d7148b5ccb5fc9f997a8ef40ff050be7092`) + - backend codegen on `sm_90, ptx94`: `succeeded` — NVVM verified and compiled an annotated kernel retaining the exact instruction on 2026-10-12. PTX SHA-256 967320791ceec33c1182d5db28170645c6977475e98c6198e8ff2ae9434f5fb6. + - device link on `sm_90, ptx94`: `succeeded` — nvJitLink13.4 linked libNVVM-produced PTX on2026-10-12. Cubin SHA-256 43d34256a5eb2a13965870356cdf8739e5d16a1036ee654f9ce51494728e089d. Tool `/usr/local/cuda/lib64/libnvJitLink.so.13.4.92` reports `V13.4.92` (SHA-256 `3afdb9677477dd584d3712a9bf2978d085729f0ed6e6432027e925d1e312e3c9`). Artifact `cubin`. - `mbarrier_arrive_expect_tx`: runtime `unexecuted` - `llvm_nvptx` uses `inline_ptx` from profile `rust-llvm-23.1.3-aab11ee0` at PTX 8.0 / sm_90+: status `validated` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) - backend codegen on `sm_90, ptx80`: `succeeded` — Revalidated 5 separate probes with nightly-2026-10-11 (rustc 6f5da406b56c2db6607202b63daee73aa8c61297; LLVM aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019): the exact reviewed PTX instruction shapes passed. The sorted ID/artifact-kind/SHA-256 ledger in refresh-2026-10-11/stages.jsonl has SHA-256 43181bf08ce9b3b668f72d19f5c2ff1fda31995dbff3b3e517bbe3ede29e620c. @@ -10013,6 +10048,13 @@ This table is generated from the resolved catalog. Evidence stages below disting - `lib_nvvm` uses `inline_ptx` from profile `cuda-13.3-libnvvm-13.3.33-mbarrier-extended` at PTX 8.0 / sm_90+: status `validated` (`CUDA toolkit 13.3; nvcc/libNVVM release V13.3.33`, SHA-256 `b6a5db233133fb0ed693f8cb8aeb1d405254ce0ab10364478985292ec98cfcf0`) - backend codegen on `compute_90, ptx93`: `succeeded` — CUDA 13.3 libNVVM verified the modules and retained each exact extended-mbarrier or async-proxy fence instruction. - device link on `sm_90, ptx93`: `succeeded` — CUDA 13.3 nvJitLink linked the libNVVM-produced LTOIR into cubins for sm_90. Tool `/usr/local/cuda-13.3/lib64/libnvJitLink.so.13.3.33` reports `V13.3.33` (SHA-256 `3ba1e744347cd68617b862eccfd98b125482e882b7a6319f42abc9a768513db8`). Artifact `cubin`. +- `mbarrier_try_wait_parity_cluster_relaxed`: runtime `unexecuted` + - `llvm_nvptx` uses `inline_ptx` from profile `rust-llvm-23.1.3-aab11ee0` at PTX 8.6 / sm_90+: status `validated` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) + - backend codegen on `sm_90, ptx86`: `succeeded` — Compiled and checked the exact instruction on 2026-10-12. PTX SHA-256 5dfa827aa34d633b208aea6582c07db167580a77728889722865433b820756fe. + - PTX assembly on `sm_90, ptx86`: `succeeded` — ptxas -arch=sm_90 -c accepted the generated PTX on 2026-10-12. Cubin SHA-256 5b48f889f189710480c0d4fd90299a85be5b7a992a14946f0328ab5f4e0ad163. Tool `/usr/local/cuda/bin/ptxas` reports `Cuda compilation tools, release 13.4, V13.4.92` (SHA-256 `6f6a701573ac9f5c2409708b3b700cb9b2a0f11b13c20468153c13b597b2d6d3`). Artifact `cubin`. + - `lib_nvvm` uses `inline_ptx` from profile `cuda-13.4-libnvvm-13.4.92-mbarrier-relaxed-cluster` at PTX 8.6 / sm_90+: status `validated` (`CUDA toolkit 13.4; nvcc/libNVVM release V13.4.92`, SHA-256 `eeef1ca79dc4ded9cc4ceb69a12f6d7148b5ccb5fc9f997a8ef40ff050be7092`) + - backend codegen on `sm_90, ptx94`: `succeeded` — NVVM verified and compiled an annotated kernel retaining the exact instruction on 2026-10-12. PTX SHA-256 455a80dd009512b448d952113f7fe899b5921f3aab1dafc256e7adc0d61098b1. + - device link on `sm_90, ptx94`: `succeeded` — nvJitLink13.4 linked libNVVM-produced PTX on2026-10-12. Cubin SHA-256 0776be219a56271a716559b413dcf4ca826e561812e72d9b98c3d0f1e4c1153c. Tool `/usr/local/cuda/lib64/libnvJitLink.so.13.4.92` reports `V13.4.92` (SHA-256 `3afdb9677477dd584d3712a9bf2978d085729f0ed6e6432027e925d1e312e3c9`). Artifact `cubin`. - `min_bf16`: runtime `not recorded` - `llvm_nvptx` uses `inline_ptx` from profile `rust-llvm-23.1.3-aab11ee0` at PTX 7.0 / sm_80+: status `validated` (`LLVM version 23.1.3-rust-1.101.0-nightly`, SHA-256 `7534b0236b9b26c7e6910548038b2285ee1bb5908e4700764dc78a7a1a54b449`) - declaration canonicalization on `sm_80, ptx70`: `succeeded` — Revalidated 12 separate probes with nightly-2026-10-11 (rustc 6f5da406b56c2db6607202b63daee73aa8c61297; LLVM aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019): canonical typed declarations and their exact imported attributes passed. The sorted ID/artifact-kind/SHA-256 ledger in refresh-2026-10-11/stages.jsonl has SHA-256 fdae9fff36f881da4949ef1733359a4d7a42da97aef5b91c6fcc54110e92fdda. diff --git a/cuda-oxide/intrinsics/overlay/mbarrier_extended.toml b/cuda-oxide/intrinsics/overlay/mbarrier_extended.toml index a314be9244..82e7ed7887 100644 --- a/cuda-oxide/intrinsics/overlay/mbarrier_extended.toml +++ b/cuda-oxide/intrinsics/overlay/mbarrier_extended.toml @@ -6,6 +6,7 @@ family = "mbarrier_extended" [mbarrier_extended] llvm_evidence_profile = "rust-llvm-23.1.3-aab11ee0" libnvvm_evidence_profile = "cuda-13.3-libnvvm-13.3.33-mbarrier-extended" +relaxed_libnvvm_evidence_profile = "cuda-13.4-libnvvm-13.4.92-mbarrier-relaxed-cluster" runtime_validation = "unexecuted" [[mbarrier_extended.variant]] @@ -51,3 +52,15 @@ operation = "fence_proxy_async_generic_acquire_shared_cluster_cluster" [[mbarrier_extended.variant]] abi_id = "i0316" operation = "nanosleep" + +[[mbarrier_extended.variant]] +abi_id = "i1031" +operation = "arrive_remote_cluster_relaxed" + +[[mbarrier_extended.variant]] +abi_id = "i1032" +operation = "try_wait_parity_cluster_relaxed" + +[[mbarrier_extended.variant]] +abi_id = "i1033" +operation = "fence_acquire_shared_cluster_cluster" diff --git a/cuda-oxide/intrinsics/probes/abs_bf16x2.ll b/cuda-oxide/intrinsics/probes/abs_bf16x2.ll index 8d4dc5ff9b..4c62d1c88f 100644 --- a/cuda-oxide/intrinsics/probes/abs_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/abs_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/abs_f16x2.ll b/cuda-oxide/intrinsics/probes/abs_f16x2.ll index 6b0ba3169e..81357e6150 100644 --- a/cuda-oxide/intrinsics/probes/abs_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/abs_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/active_mask.ll b/cuda-oxide/intrinsics/probes/active_mask.ll index 4e8e1cb75b..d217c76860 100644 --- a/cuda-oxide/intrinsics/probes/active_mask.ll +++ b/cuda-oxide/intrinsics/probes/active_mask.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_bf16x2.ll b/cuda-oxide/intrinsics/probes/add_bf16x2.ll index a16065c8ed..23c78be416 100644 --- a/cuda-oxide/intrinsics/probes/add_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/add_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_f16x2.ll b/cuda-oxide/intrinsics/probes/add_f16x2.ll index 5f47860b12..9511004d23 100644 --- a/cuda-oxide/intrinsics/probes/add_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/add_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_f32x2.ll b/cuda-oxide/intrinsics/probes/add_f32x2.ll index 2d7b242f7e..a044af62e4 100644 --- a/cuda-oxide/intrinsics/probes/add_f32x2.ll +++ b/cuda-oxide/intrinsics/probes/add_f32x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_ftz_f32x2.ll b/cuda-oxide/intrinsics/probes/add_ftz_f32x2.ll index a9178bf880..dbe43d4a28 100644 --- a/cuda-oxide/intrinsics/probes/add_ftz_f32x2.ll +++ b/cuda-oxide/intrinsics/probes/add_ftz_f32x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rm_f32.ll b/cuda-oxide/intrinsics/probes/add_rm_f32.ll index d9caf29c4d..cc8b325a6c 100644 --- a/cuda-oxide/intrinsics/probes/add_rm_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rm_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rm_f64.ll b/cuda-oxide/intrinsics/probes/add_rm_f64.ll index d6427adfec..ab89d0a214 100644 --- a/cuda-oxide/intrinsics/probes/add_rm_f64.ll +++ b/cuda-oxide/intrinsics/probes/add_rm_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rm_ftz_f32.ll b/cuda-oxide/intrinsics/probes/add_rm_ftz_f32.ll index b310e10aa6..caa73de942 100644 --- a/cuda-oxide/intrinsics/probes/add_rm_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rm_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rm_ftz_sat_f32.ll b/cuda-oxide/intrinsics/probes/add_rm_ftz_sat_f32.ll index 0fdc8fecac..ed022e1a60 100644 --- a/cuda-oxide/intrinsics/probes/add_rm_ftz_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rm_ftz_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rm_sat_f32.ll b/cuda-oxide/intrinsics/probes/add_rm_sat_f32.ll index 1d5b5a8d17..d7ea2a6c28 100644 --- a/cuda-oxide/intrinsics/probes/add_rm_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rm_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rn_f32.ll b/cuda-oxide/intrinsics/probes/add_rn_f32.ll index ac478e443a..ca8a68ec65 100644 --- a/cuda-oxide/intrinsics/probes/add_rn_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rn_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rn_f64.ll b/cuda-oxide/intrinsics/probes/add_rn_f64.ll index 868718c07d..cf3854a657 100644 --- a/cuda-oxide/intrinsics/probes/add_rn_f64.ll +++ b/cuda-oxide/intrinsics/probes/add_rn_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rn_ftz_f32.ll b/cuda-oxide/intrinsics/probes/add_rn_ftz_f32.ll index a3b949d3aa..9c6667814d 100644 --- a/cuda-oxide/intrinsics/probes/add_rn_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rn_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rn_ftz_sat_f32.ll b/cuda-oxide/intrinsics/probes/add_rn_ftz_sat_f32.ll index 5687dd1101..3580e66348 100644 --- a/cuda-oxide/intrinsics/probes/add_rn_ftz_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rn_ftz_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rn_sat_f32.ll b/cuda-oxide/intrinsics/probes/add_rn_sat_f32.ll index bb58718e77..a9d55d8e3d 100644 --- a/cuda-oxide/intrinsics/probes/add_rn_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rn_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rp_f32.ll b/cuda-oxide/intrinsics/probes/add_rp_f32.ll index 79be1c6d39..eedd481f5f 100644 --- a/cuda-oxide/intrinsics/probes/add_rp_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rp_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rp_f64.ll b/cuda-oxide/intrinsics/probes/add_rp_f64.ll index dca64b9d6f..30fc335de0 100644 --- a/cuda-oxide/intrinsics/probes/add_rp_f64.ll +++ b/cuda-oxide/intrinsics/probes/add_rp_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rp_ftz_f32.ll b/cuda-oxide/intrinsics/probes/add_rp_ftz_f32.ll index 3176078a27..0544123aab 100644 --- a/cuda-oxide/intrinsics/probes/add_rp_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rp_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rp_ftz_sat_f32.ll b/cuda-oxide/intrinsics/probes/add_rp_ftz_sat_f32.ll index f99a3498b5..ee91f565c7 100644 --- a/cuda-oxide/intrinsics/probes/add_rp_ftz_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rp_ftz_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rp_sat_f32.ll b/cuda-oxide/intrinsics/probes/add_rp_sat_f32.ll index 4dd40c4bd4..51b7ab41d2 100644 --- a/cuda-oxide/intrinsics/probes/add_rp_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rp_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rz_f32.ll b/cuda-oxide/intrinsics/probes/add_rz_f32.ll index 5222452dae..0e9dbc6378 100644 --- a/cuda-oxide/intrinsics/probes/add_rz_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rz_f64.ll b/cuda-oxide/intrinsics/probes/add_rz_f64.ll index 1b0f4fa957..f2ec47fa9d 100644 --- a/cuda-oxide/intrinsics/probes/add_rz_f64.ll +++ b/cuda-oxide/intrinsics/probes/add_rz_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rz_ftz_f32.ll b/cuda-oxide/intrinsics/probes/add_rz_ftz_f32.ll index efa1aca07f..996ef0f819 100644 --- a/cuda-oxide/intrinsics/probes/add_rz_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rz_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rz_ftz_sat_f32.ll b/cuda-oxide/intrinsics/probes/add_rz_ftz_sat_f32.ll index ba2f314e8f..5da7787ce3 100644 --- a/cuda-oxide/intrinsics/probes/add_rz_ftz_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rz_ftz_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/add_rz_sat_f32.ll b/cuda-oxide/intrinsics/probes/add_rz_sat_f32.ll index 4addff21bf..5617409edb 100644 --- a/cuda-oxide/intrinsics/probes/add_rz_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/add_rz_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/all_sync.ll b/cuda-oxide/intrinsics/probes/all_sync.ll index 06f6bca606..0d7b7ddd89 100644 --- a/cuda-oxide/intrinsics/probes/all_sync.ll +++ b/cuda-oxide/intrinsics/probes/all_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/any_sync.ll b/cuda-oxide/intrinsics/probes/any_sync.ll index 9e64411b59..61de9a2078 100644 --- a/cuda-oxide/intrinsics/probes/any_sync.ll +++ b/cuda-oxide/intrinsics/probes/any_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ballot_sync.ll b/cuda-oxide/intrinsics/probes/ballot_sync.ll index 693ceb4988..eec7282c31 100644 --- a/cuda-oxide/intrinsics/probes/ballot_sync.ll +++ b/cuda-oxide/intrinsics/probes/ballot_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/barrier_cluster_arrive.ll b/cuda-oxide/intrinsics/probes/barrier_cluster_arrive.ll index 0df822aab0..bd4d344896 100644 --- a/cuda-oxide/intrinsics/probes/barrier_cluster_arrive.ll +++ b/cuda-oxide/intrinsics/probes/barrier_cluster_arrive.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/barrier_cluster_arrive_aligned.ll b/cuda-oxide/intrinsics/probes/barrier_cluster_arrive_aligned.ll index a34bb6cfd2..cec70e9b1f 100644 --- a/cuda-oxide/intrinsics/probes/barrier_cluster_arrive_aligned.ll +++ b/cuda-oxide/intrinsics/probes/barrier_cluster_arrive_aligned.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/barrier_cluster_arrive_relaxed.ll b/cuda-oxide/intrinsics/probes/barrier_cluster_arrive_relaxed.ll index 367c8e7576..d9a2a00ba8 100644 --- a/cuda-oxide/intrinsics/probes/barrier_cluster_arrive_relaxed.ll +++ b/cuda-oxide/intrinsics/probes/barrier_cluster_arrive_relaxed.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/barrier_cluster_arrive_relaxed_aligned.ll b/cuda-oxide/intrinsics/probes/barrier_cluster_arrive_relaxed_aligned.ll index f81ce1e7c1..97cfe0de0c 100644 --- a/cuda-oxide/intrinsics/probes/barrier_cluster_arrive_relaxed_aligned.ll +++ b/cuda-oxide/intrinsics/probes/barrier_cluster_arrive_relaxed_aligned.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/barrier_cluster_wait.ll b/cuda-oxide/intrinsics/probes/barrier_cluster_wait.ll index 1630c0bc89..7f9c225763 100644 --- a/cuda-oxide/intrinsics/probes/barrier_cluster_wait.ll +++ b/cuda-oxide/intrinsics/probes/barrier_cluster_wait.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/barrier_cluster_wait_aligned.ll b/cuda-oxide/intrinsics/probes/barrier_cluster_wait_aligned.ll index aeff105121..8f93f13ce4 100644 --- a/cuda-oxide/intrinsics/probes/barrier_cluster_wait_aligned.ll +++ b/cuda-oxide/intrinsics/probes/barrier_cluster_wait_aligned.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/barrier_cta_arrive.ll b/cuda-oxide/intrinsics/probes/barrier_cta_arrive.ll index 5239ad7c70..4940853a10 100644 --- a/cuda-oxide/intrinsics/probes/barrier_cta_arrive.ll +++ b/cuda-oxide/intrinsics/probes/barrier_cta_arrive.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/barrier_cta_arrive_aligned.ll b/cuda-oxide/intrinsics/probes/barrier_cta_arrive_aligned.ll index 21337f2cc7..e760d35930 100644 --- a/cuda-oxide/intrinsics/probes/barrier_cta_arrive_aligned.ll +++ b/cuda-oxide/intrinsics/probes/barrier_cta_arrive_aligned.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/barrier_cta_sync.ll b/cuda-oxide/intrinsics/probes/barrier_cta_sync.ll index fdd6c6a3c7..6f9228534c 100644 --- a/cuda-oxide/intrinsics/probes/barrier_cta_sync.ll +++ b/cuda-oxide/intrinsics/probes/barrier_cta_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/barrier_cta_sync_aligned.ll b/cuda-oxide/intrinsics/probes/barrier_cta_sync_aligned.ll index d6592c61f3..b6ef2961d2 100644 --- a/cuda-oxide/intrinsics/probes/barrier_cta_sync_aligned.ll +++ b/cuda-oxide/intrinsics/probes/barrier_cta_sync_aligned.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/block_dim_x.ll b/cuda-oxide/intrinsics/probes/block_dim_x.ll index e67011acb8..10b4308e05 100644 --- a/cuda-oxide/intrinsics/probes/block_dim_x.ll +++ b/cuda-oxide/intrinsics/probes/block_dim_x.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/block_dim_y.ll b/cuda-oxide/intrinsics/probes/block_dim_y.ll index 786b3bb6c8..65cbc78d96 100644 --- a/cuda-oxide/intrinsics/probes/block_dim_y.ll +++ b/cuda-oxide/intrinsics/probes/block_dim_y.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/block_dim_z.ll b/cuda-oxide/intrinsics/probes/block_dim_z.ll index 28af10ba4d..6a75ceee59 100644 --- a/cuda-oxide/intrinsics/probes/block_dim_z.ll +++ b/cuda-oxide/intrinsics/probes/block_dim_z.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/block_idx_x.ll b/cuda-oxide/intrinsics/probes/block_idx_x.ll index 9a065ad112..1b23798f62 100644 --- a/cuda-oxide/intrinsics/probes/block_idx_x.ll +++ b/cuda-oxide/intrinsics/probes/block_idx_x.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/block_idx_y.ll b/cuda-oxide/intrinsics/probes/block_idx_y.ll index bda2b6470e..6b0c49a99e 100644 --- a/cuda-oxide/intrinsics/probes/block_idx_y.ll +++ b/cuda-oxide/intrinsics/probes/block_idx_y.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/block_idx_z.ll b/cuda-oxide/intrinsics/probes/block_idx_z.ll index 3880c5a102..2ed13c669c 100644 --- a/cuda-oxide/intrinsics/probes/block_idx_z.ll +++ b/cuda-oxide/intrinsics/probes/block_idx_z.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/breakpoint.ll b/cuda-oxide/intrinsics/probes/breakpoint.ll index 46dee7a14d..5bdcad0091 100644 --- a/cuda-oxide/intrinsics/probes/breakpoint.ll +++ b/cuda-oxide/intrinsics/probes/breakpoint.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/clc_query_get_first_ctaid_x.ll b/cuda-oxide/intrinsics/probes/clc_query_get_first_ctaid_x.ll index cf6157f5e7..fb09465607 100644 --- a/cuda-oxide/intrinsics/probes/clc_query_get_first_ctaid_x.ll +++ b/cuda-oxide/intrinsics/probes/clc_query_get_first_ctaid_x.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/clc_query_get_first_ctaid_y.ll b/cuda-oxide/intrinsics/probes/clc_query_get_first_ctaid_y.ll index b1836c955c..28287205ed 100644 --- a/cuda-oxide/intrinsics/probes/clc_query_get_first_ctaid_y.ll +++ b/cuda-oxide/intrinsics/probes/clc_query_get_first_ctaid_y.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/clc_query_get_first_ctaid_z.ll b/cuda-oxide/intrinsics/probes/clc_query_get_first_ctaid_z.ll index f190856ae2..8c99ab4d5e 100644 --- a/cuda-oxide/intrinsics/probes/clc_query_get_first_ctaid_z.ll +++ b/cuda-oxide/intrinsics/probes/clc_query_get_first_ctaid_z.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/clc_query_is_canceled.ll b/cuda-oxide/intrinsics/probes/clc_query_is_canceled.ll index 799b7f6664..3b19545be5 100644 --- a/cuda-oxide/intrinsics/probes/clc_query_is_canceled.ll +++ b/cuda-oxide/intrinsics/probes/clc_query_is_canceled.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/clc_try_cancel.ll b/cuda-oxide/intrinsics/probes/clc_try_cancel.ll index 18a5df02c4..c9b8603b14 100644 --- a/cuda-oxide/intrinsics/probes/clc_try_cancel.ll +++ b/cuda-oxide/intrinsics/probes/clc_try_cancel.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/clc_try_cancel_multicast.ll b/cuda-oxide/intrinsics/probes/clc_try_cancel_multicast.ll index 607e92e85a..4bd10293c1 100644 --- a/cuda-oxide/intrinsics/probes/clc_try_cancel_multicast.ll +++ b/cuda-oxide/intrinsics/probes/clc_try_cancel_multicast.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/clock.lib_nvvm.ll b/cuda-oxide/intrinsics/probes/clock.lib_nvvm.ll index 4bc836ab4f..7ae16ba379 100644 --- a/cuda-oxide/intrinsics/probes/clock.lib_nvvm.ll +++ b/cuda-oxide/intrinsics/probes/clock.lib_nvvm.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/clock.llvm_nvptx.ll b/cuda-oxide/intrinsics/probes/clock.llvm_nvptx.ll index d73af18975..2873b62fcc 100644 --- a/cuda-oxide/intrinsics/probes/clock.llvm_nvptx.ll +++ b/cuda-oxide/intrinsics/probes/clock.llvm_nvptx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/clock64.lib_nvvm.ll b/cuda-oxide/intrinsics/probes/clock64.lib_nvvm.ll index c913d437cd..ec4d703e22 100644 --- a/cuda-oxide/intrinsics/probes/clock64.lib_nvvm.ll +++ b/cuda-oxide/intrinsics/probes/clock64.lib_nvvm.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/clock64.llvm_nvptx.ll b/cuda-oxide/intrinsics/probes/clock64.llvm_nvptx.ll index 88e0fea129..f503d8f80b 100644 --- a/cuda-oxide/intrinsics/probes/clock64.llvm_nvptx.ll +++ b/cuda-oxide/intrinsics/probes/clock64.llvm_nvptx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_block_count.ll b/cuda-oxide/intrinsics/probes/cluster_block_count.ll index 6d67f4006c..47c54dce48 100644 --- a/cuda-oxide/intrinsics/probes/cluster_block_count.ll +++ b/cuda-oxide/intrinsics/probes/cluster_block_count.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_block_idx_x.ll b/cuda-oxide/intrinsics/probes/cluster_block_idx_x.ll index 946fe47af9..247757bb23 100644 --- a/cuda-oxide/intrinsics/probes/cluster_block_idx_x.ll +++ b/cuda-oxide/intrinsics/probes/cluster_block_idx_x.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_block_idx_y.ll b/cuda-oxide/intrinsics/probes/cluster_block_idx_y.ll index 2cf7655aac..b6c102450a 100644 --- a/cuda-oxide/intrinsics/probes/cluster_block_idx_y.ll +++ b/cuda-oxide/intrinsics/probes/cluster_block_idx_y.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_block_idx_z.ll b/cuda-oxide/intrinsics/probes/cluster_block_idx_z.ll index d60188aeea..7248e04c1c 100644 --- a/cuda-oxide/intrinsics/probes/cluster_block_idx_z.ll +++ b/cuda-oxide/intrinsics/probes/cluster_block_idx_z.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_block_rank.ll b/cuda-oxide/intrinsics/probes/cluster_block_rank.ll index 86ce5bd629..f257313eb2 100644 --- a/cuda-oxide/intrinsics/probes/cluster_block_rank.ll +++ b/cuda-oxide/intrinsics/probes/cluster_block_rank.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_dim_x.ll b/cuda-oxide/intrinsics/probes/cluster_dim_x.ll index 8b2fab7181..8857c29f83 100644 --- a/cuda-oxide/intrinsics/probes/cluster_dim_x.ll +++ b/cuda-oxide/intrinsics/probes/cluster_dim_x.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_dim_y.ll b/cuda-oxide/intrinsics/probes/cluster_dim_y.ll index a9edb71987..c08b523853 100644 --- a/cuda-oxide/intrinsics/probes/cluster_dim_y.ll +++ b/cuda-oxide/intrinsics/probes/cluster_dim_y.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_dim_z.ll b/cuda-oxide/intrinsics/probes/cluster_dim_z.ll index 6c56d59739..c9ed85f268 100644 --- a/cuda-oxide/intrinsics/probes/cluster_dim_z.ll +++ b/cuda-oxide/intrinsics/probes/cluster_dim_z.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_grid_dim_x.ll b/cuda-oxide/intrinsics/probes/cluster_grid_dim_x.ll index 7a0d0f2d1f..428a0aaef4 100644 --- a/cuda-oxide/intrinsics/probes/cluster_grid_dim_x.ll +++ b/cuda-oxide/intrinsics/probes/cluster_grid_dim_x.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_grid_dim_y.ll b/cuda-oxide/intrinsics/probes/cluster_grid_dim_y.ll index 501c5ea94d..44319cac56 100644 --- a/cuda-oxide/intrinsics/probes/cluster_grid_dim_y.ll +++ b/cuda-oxide/intrinsics/probes/cluster_grid_dim_y.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_grid_dim_z.ll b/cuda-oxide/intrinsics/probes/cluster_grid_dim_z.ll index ffd6c030a5..5007979ba5 100644 --- a/cuda-oxide/intrinsics/probes/cluster_grid_dim_z.ll +++ b/cuda-oxide/intrinsics/probes/cluster_grid_dim_z.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_idx_x.ll b/cuda-oxide/intrinsics/probes/cluster_idx_x.ll index bd60cfe1c4..4905ff632b 100644 --- a/cuda-oxide/intrinsics/probes/cluster_idx_x.ll +++ b/cuda-oxide/intrinsics/probes/cluster_idx_x.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_idx_y.ll b/cuda-oxide/intrinsics/probes/cluster_idx_y.ll index aabf5a66af..3bb24eef4d 100644 --- a/cuda-oxide/intrinsics/probes/cluster_idx_y.ll +++ b/cuda-oxide/intrinsics/probes/cluster_idx_y.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cluster_idx_z.ll b/cuda-oxide/intrinsics/probes/cluster_idx_z.ll index b490a9f1d6..9a1ebf771c 100644 --- a/cuda-oxide/intrinsics/probes/cluster_idx_z.ll +++ b/cuda-oxide/intrinsics/probes/cluster_idx_z.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cos_approx_f32.ll b/cuda-oxide/intrinsics/probes/cos_approx_f32.ll index 1784bb6834..49bf0f843d 100644 --- a/cuda-oxide/intrinsics/probes/cos_approx_f32.ll +++ b/cuda-oxide/intrinsics/probes/cos_approx_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cos_approx_ftz_f32.ll b/cuda-oxide/intrinsics/probes/cos_approx_ftz_f32.ll index 87c50d6f0c..ec5ad2066c 100644 --- a/cuda-oxide/intrinsics/probes/cos_approx_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/cos_approx_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_commit_group.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_commit_group.ll index 75ab141171..f78dd4b261 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_commit_group.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_commit_group.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_1d_l2.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_1d_l2.ll index 98b38adf23..169f6afc96 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_1d_l2.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_1d_l2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_1d_l2_cache_hint.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_1d_l2_cache_hint.ll index 3dc3da8fcd..6b54210bdf 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_1d_l2_cache_hint.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_1d_l2_cache_hint.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_2d_l2.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_2d_l2.ll index c04c886de7..73ed39bfdb 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_2d_l2.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_2d_l2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_2d_l2_cache_hint.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_2d_l2_cache_hint.ll index 1400907b3b..1550bdb146 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_2d_l2_cache_hint.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_2d_l2_cache_hint.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_3d_l2.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_3d_l2.ll index df303a21fe..5989c77572 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_3d_l2.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_3d_l2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_3d_l2_cache_hint.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_3d_l2_cache_hint.ll index fd07ecc1c5..3558e036d2 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_3d_l2_cache_hint.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_3d_l2_cache_hint.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_4d_l2.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_4d_l2.ll index 6499be99d2..4ff7a7a30e 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_4d_l2.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_4d_l2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_4d_l2_cache_hint.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_4d_l2_cache_hint.ll index 92aa106a4c..73e29289d2 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_4d_l2_cache_hint.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_4d_l2_cache_hint.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_5d_l2.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_5d_l2.ll index 2f6bb6b969..a8ebbdc961 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_5d_l2.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_5d_l2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_5d_l2_cache_hint.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_5d_l2_cache_hint.ll index 33fdeeb4ef..3fd5be56bd 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_5d_l2_cache_hint.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_5d_l2_cache_hint.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_gather4_2d_l2.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_gather4_2d_l2.ll index 8f8245581d..f1f1bef3a0 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_gather4_2d_l2.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_gather4_2d_l2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint.ll index 79a89f1aec..7a5fa37083 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_1d_g2s.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_1d_g2s.ll index 7653e1fea8..e116c77069 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_1d_g2s.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_1d_g2s.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_1d_s2g.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_1d_s2g.ll index aaabeacb00..d2a7ec2005 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_1d_s2g.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_1d_s2g.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_g2s.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_g2s.ll index 404677b926..276b7e9163 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_g2s.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_g2s.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_g2s_multicast.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_g2s_multicast.ll index 2f73720d47..ffd24a44e3 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_g2s_multicast.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_g2s_multicast.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_g2s_multicast_cg2.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_g2s_multicast_cg2.ll index 034e0a2d7e..130bac8dbd 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_g2s_multicast_cg2.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_g2s_multicast_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_s2g.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_s2g.ll index 3f5c55d733..55651796e6 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_s2g.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_2d_s2g.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_3d_g2s.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_3d_g2s.ll index 815524e13d..23b6965183 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_3d_g2s.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_3d_g2s.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_3d_s2g.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_3d_s2g.ll index 5f4c68de86..473f0bcf05 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_3d_s2g.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_3d_s2g.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_4d_g2s.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_4d_g2s.ll index fb315ab17b..8af10fa332 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_4d_g2s.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_4d_g2s.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_4d_s2g.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_4d_s2g.ll index b34a359810..102ca6e2f1 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_4d_s2g.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_4d_s2g.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_5d_g2s.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_5d_g2s.ll index c944cbc122..faada17a6e 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_5d_g2s.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_5d_g2s.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_5d_s2g.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_5d_s2g.ll index 3d41d8dd46..259f1be65b 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_5d_s2g.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_5d_s2g.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_im2col_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_im2col_3d.ll index 3cfe85fa21..fd3c5688da 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_im2col_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_im2col_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_im2col_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_im2col_4d.ll index ed9a243e15..14ff40d703 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_im2col_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_im2col_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_im2col_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_im2col_5d.ll index 1f011de2bf..2b40d09a48 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_im2col_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_im2col_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_1d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_1d.ll index e8fd0ebdeb..0206454f84 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_1d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_1d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_2d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_2d.ll index d9b39e2dd4..9855e29662 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_2d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_2d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_3d.ll index da13f1c323..1397756bec 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_4d.ll index b2474cd7bb..e90f582b35 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_5d.ll index d1861b15db..2d388a8ce8 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_add_tile_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_im2col_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_im2col_3d.ll index 407f788b4e..818db0699a 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_im2col_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_im2col_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_im2col_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_im2col_4d.ll index 102038dc76..a70d7fbc26 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_im2col_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_im2col_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_im2col_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_im2col_5d.ll index a294409761..7aa2946634 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_im2col_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_im2col_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_1d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_1d.ll index 9f0f1357dc..f182eefd23 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_1d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_1d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_2d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_2d.ll index ff1f2382de..954df0b8e8 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_2d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_2d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_3d.ll index 1101a918ee..d759fac550 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_4d.ll index cc55c234f9..dd6c45c696 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_5d.ll index c21304669d..c4500dc8ff 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_and_tile_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_im2col_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_im2col_3d.ll index 781637fa14..a5484e5cd1 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_im2col_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_im2col_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_im2col_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_im2col_4d.ll index 98ee7dd3b7..f97e69cd65 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_im2col_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_im2col_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_im2col_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_im2col_5d.ll index 3eeded37bb..2d3ed3b163 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_im2col_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_im2col_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_1d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_1d.ll index 14602c66a8..b4c47482c1 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_1d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_1d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_2d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_2d.ll index 5c4e3aa5dd..ed59b284fd 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_2d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_2d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_3d.ll index bf18d02a3e..536c01e740 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_4d.ll index 50d9fcc6da..301e8fbf46 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_5d.ll index af3b448e4e..576a768477 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_dec_tile_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_im2col_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_im2col_3d.ll index 9be1d515a0..3347f3fef7 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_im2col_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_im2col_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_im2col_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_im2col_4d.ll index 1f18d79e37..1071d64a43 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_im2col_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_im2col_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_im2col_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_im2col_5d.ll index ae69b28613..15adc8c67b 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_im2col_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_im2col_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_1d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_1d.ll index d2b60d118a..5ee18efb71 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_1d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_1d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_2d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_2d.ll index 838225b14e..7134a9184a 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_2d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_2d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_3d.ll index 264a46350b..626c25fb92 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_4d.ll index d6f4761178..91ef231d03 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_5d.ll index 6a88d31410..4c0d153435 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_inc_tile_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_im2col_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_im2col_3d.ll index 801bed77c6..18a517347e 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_im2col_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_im2col_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_im2col_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_im2col_4d.ll index f81e23096c..757cc2e741 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_im2col_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_im2col_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_im2col_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_im2col_5d.ll index bea78ec511..c10761b6dc 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_im2col_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_im2col_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_1d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_1d.ll index 05e156cf21..9e38a48a2b 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_1d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_1d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_2d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_2d.ll index 2868049869..35156ebe8d 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_2d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_2d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_3d.ll index 62ca1d3c46..580d2781a7 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_4d.ll index a067fbc432..4c563c212b 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_5d.ll index 86210f541c..c8136e688b 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_max_tile_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_im2col_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_im2col_3d.ll index 9185e14265..a49058b908 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_im2col_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_im2col_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_im2col_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_im2col_4d.ll index ef72aacc98..27bfa46daa 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_im2col_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_im2col_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_im2col_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_im2col_5d.ll index 4252111261..f77d57f5c4 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_im2col_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_im2col_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_1d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_1d.ll index b1161632c5..1a31daee9b 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_1d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_1d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_2d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_2d.ll index 0748453edc..77a29af478 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_2d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_2d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_3d.ll index 04eae963d8..ab44c58060 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_4d.ll index 338781c03c..f7adec88ad 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_5d.ll index 736bb224a0..c3f6014d82 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_min_tile_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_im2col_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_im2col_3d.ll index 0630f594c9..4c9b1dc416 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_im2col_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_im2col_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_im2col_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_im2col_4d.ll index 7d875bf17f..e905bfadb4 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_im2col_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_im2col_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_im2col_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_im2col_5d.ll index 69f5c54464..5ea01caa64 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_im2col_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_im2col_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_1d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_1d.ll index 20e3ea95d8..8580a3827e 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_1d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_1d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_2d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_2d.ll index 609ebcca7d..262e354f05 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_2d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_2d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_3d.ll index 14ac0c2eab..1ba1dab572 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_4d.ll index 6be933d1b8..e42594dd1c 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_5d.ll index beb37633d2..ed90b5484a 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_or_tile_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_im2col_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_im2col_3d.ll index ed3742231b..80589e6877 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_im2col_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_im2col_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_im2col_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_im2col_4d.ll index b7a255034e..7dbcc9029c 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_im2col_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_im2col_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_im2col_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_im2col_5d.ll index 8cde8c999b..fedf78415f 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_im2col_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_im2col_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_1d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_1d.ll index 600731f905..69e77bd4ab 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_1d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_1d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_2d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_2d.ll index bf12de8986..7767dbb6d4 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_2d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_2d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_3d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_3d.ll index 132aa3a658..09cd40c92d 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_3d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_3d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_4d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_4d.ll index 152c9b490a..dd76e74f9c 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_4d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_4d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_5d.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_5d.ll index 4984dd45c0..dafed42092 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_5d.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_tensor_reduce_xor_tile_5d.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_wait_group.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_wait_group.ll index c62366704e..aef22d3a84 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_wait_group.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_wait_group.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_bulk_wait_group_read.ll b/cuda-oxide/intrinsics/probes/cp_async_bulk_wait_group_read.ll index 7904df25e2..24c72db8e8 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_bulk_wait_group_read.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_bulk_wait_group_read.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_ca_16.ll b/cuda-oxide/intrinsics/probes/cp_async_ca_16.ll index 7487a5c136..3eda482fc8 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_ca_16.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_ca_16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_ca_4.ll b/cuda-oxide/intrinsics/probes/cp_async_ca_4.ll index 62f92fd92f..308c477558 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_ca_4.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_ca_4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_ca_8.ll b/cuda-oxide/intrinsics/probes/cp_async_ca_8.ll index 43769e2c41..a1bfa9d528 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_ca_8.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_ca_8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_ca_zfill_16.ll b/cuda-oxide/intrinsics/probes/cp_async_ca_zfill_16.ll index 29e172b0cf..ce2349c1f7 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_ca_zfill_16.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_ca_zfill_16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_ca_zfill_4.ll b/cuda-oxide/intrinsics/probes/cp_async_ca_zfill_4.ll index fd58644ea6..0cfa05fa6d 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_ca_zfill_4.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_ca_zfill_4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_ca_zfill_8.ll b/cuda-oxide/intrinsics/probes/cp_async_ca_zfill_8.ll index 1f1aeadab4..8f804b400e 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_ca_zfill_8.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_ca_zfill_8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_cg_16.ll b/cuda-oxide/intrinsics/probes/cp_async_cg_16.ll index 1690212a5d..b3c872d4b4 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_cg_16.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_cg_16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_cg_zfill_16.ll b/cuda-oxide/intrinsics/probes/cp_async_cg_zfill_16.ll index d66524c7aa..c85a0a390c 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_cg_zfill_16.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_cg_zfill_16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_commit_group.ll b/cuda-oxide/intrinsics/probes/cp_async_commit_group.ll index 313ed7933a..1449e37c10 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_commit_group.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_commit_group.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive.ll b/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive.ll index 5a8dbe8acd..74fdcc7799 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive_noinc.ll b/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive_noinc.ll index 49f678ab1d..297755635b 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive_noinc.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive_noinc.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive_noinc_shared.ll b/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive_noinc_shared.ll index 7f9934c678..6b1ebc7615 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive_noinc_shared.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive_noinc_shared.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive_shared.ll b/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive_shared.ll index e968ce91d4..68d48df648 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive_shared.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_mbarrier_arrive_shared.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_wait_all.ll b/cuda-oxide/intrinsics/probes/cp_async_wait_all.ll index 7a0fe4ce14..541d2c4bcf 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_wait_all.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_wait_all.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cp_async_wait_group.ll b/cuda-oxide/intrinsics/probes/cp_async_wait_group.ll index c795a3f340..d89664d7f5 100644 --- a/cuda-oxide/intrinsics/probes/cp_async_wait_group.ll +++ b/cuda-oxide/intrinsics/probes/cp_async_wait_group.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_f16x2_f32.ll b/cuda-oxide/intrinsics/probes/cvt_f16x2_f32.ll index 8241fe9fb2..56a49588fe 100644 --- a/cuda-oxide/intrinsics/probes/cvt_f16x2_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_f16x2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_f32x2_bf16x2.ll b/cuda-oxide/intrinsics/probes/cvt_f32x2_bf16x2.ll index 586f0eb223..84a130bf47 100644 --- a/cuda-oxide/intrinsics/probes/cvt_f32x2_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/cvt_f32x2_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_f16x2_e4m3x2.ll b/cuda-oxide/intrinsics/probes/cvt_rn_f16x2_e4m3x2.ll index cf66cf4214..4b1167be5c 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_f16x2_e4m3x2.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_f16x2_e4m3x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_f16x2_e5m2x2.ll b/cuda-oxide/intrinsics/probes/cvt_rn_f16x2_e5m2x2.ll index c264065fac..c71aeb2d18 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_f16x2_e5m2x2.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_f16x2_e5m2x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_relu_bf16x2_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rn_relu_bf16x2_f32.ll index 61f22c191c..496392c1d1 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_relu_bf16x2_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_relu_bf16x2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_relu_f16x2_e4m3x2.ll b/cuda-oxide/intrinsics/probes/cvt_rn_relu_f16x2_e4m3x2.ll index 1cd9d3166d..c0a17f00ee 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_relu_f16x2_e4m3x2.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_relu_f16x2_e4m3x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_relu_f16x2_e5m2x2.ll b/cuda-oxide/intrinsics/probes/cvt_rn_relu_f16x2_e5m2x2.ll index 14b7182339..999d121721 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_relu_f16x2_e5m2x2.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_relu_f16x2_e5m2x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_relu_f16x2_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rn_relu_f16x2_f32.ll index d792727fab..bcbb814757 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_relu_f16x2_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_relu_f16x2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_relu_satfinite_tf32_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rn_relu_satfinite_tf32_f32.ll index 80e692572a..7235e6108f 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_relu_satfinite_tf32_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_relu_satfinite_tf32_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_relu_tf32_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rn_relu_tf32_f32.ll index 8f29e177f2..0e31b8ba5f 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_relu_tf32_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_relu_tf32_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e4m3x2_f16x2.ll b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e4m3x2_f16x2.ll index e60d448fda..9fefb5e028 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e4m3x2_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e4m3x2_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e4m3x2_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e4m3x2_f32.ll index 5023080530..a3dfa3e715 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e4m3x2_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e4m3x2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e5m2x2_f16x2.ll b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e5m2x2_f16x2.ll index ffbb5dfb66..010be1cb43 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e5m2x2_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e5m2x2_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e5m2x2_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e5m2x2_f32.ll index b3b82ca6bf..40ac00c072 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e5m2x2_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_e5m2x2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e4m3x2_f16x2.ll b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e4m3x2_f16x2.ll index b1de5971bd..70a68df1c9 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e4m3x2_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e4m3x2_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e4m3x2_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e4m3x2_f32.ll index 0b9db76d5d..34e3436ede 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e4m3x2_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e4m3x2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e5m2x2_f16x2.ll b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e5m2x2_f16x2.ll index 2ad9bbe606..9c7361a28e 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e5m2x2_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e5m2x2_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e5m2x2_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e5m2x2_f32.ll index b15f62f00f..b653a5bd16 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e5m2x2_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_relu_e5m2x2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_tf32_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_tf32_f32.ll index bd90a41407..44d38fff16 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_tf32_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_satfinite_tf32_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rn_tf32_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rn_tf32_f32.ll index e89c9739f3..7d6209a6f2 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rn_tf32_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rn_tf32_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rna_satfinite_tf32_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rna_satfinite_tf32_f32.ll index 4547b1ae64..07bc7fed32 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rna_satfinite_tf32_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rna_satfinite_tf32_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rna_tf32_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rna_tf32_f32.ll index a75030acf8..9028d2fe0d 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rna_tf32_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rna_tf32_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rz_bf16x2_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rz_bf16x2_f32.ll index 5543fb5f5f..ec1e0b4d21 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rz_bf16x2_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rz_bf16x2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rz_f16x2_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rz_f16x2_f32.ll index 32823de6ad..c2729f17c7 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rz_f16x2_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rz_f16x2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rz_relu_satfinite_tf32_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rz_relu_satfinite_tf32_f32.ll index d31014a2df..5fb0ba7d33 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rz_relu_satfinite_tf32_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rz_relu_satfinite_tf32_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rz_relu_tf32_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rz_relu_tf32_f32.ll index 0fc77c127b..2fc2fe539e 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rz_relu_tf32_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rz_relu_tf32_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rz_satfinite_tf32_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rz_satfinite_tf32_f32.ll index 732c4d55a1..9a074b800a 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rz_satfinite_tf32_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rz_satfinite_tf32_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/cvt_rz_tf32_f32.ll b/cuda-oxide/intrinsics/probes/cvt_rz_tf32_f32.ll index ed114cc08e..e218bce136 100644 --- a/cuda-oxide/intrinsics/probes/cvt_rz_tf32_f32.ll +++ b/cuda-oxide/intrinsics/probes/cvt_rz_tf32_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/div_rm_f32.ll b/cuda-oxide/intrinsics/probes/div_rm_f32.ll index 1a466af647..8ae59725b4 100644 --- a/cuda-oxide/intrinsics/probes/div_rm_f32.ll +++ b/cuda-oxide/intrinsics/probes/div_rm_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/div_rm_f64.ll b/cuda-oxide/intrinsics/probes/div_rm_f64.ll index 3ccfe9b690..944e78c3ce 100644 --- a/cuda-oxide/intrinsics/probes/div_rm_f64.ll +++ b/cuda-oxide/intrinsics/probes/div_rm_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/div_rm_ftz_f32.ll b/cuda-oxide/intrinsics/probes/div_rm_ftz_f32.ll index e9d05f498c..93fa0404ea 100644 --- a/cuda-oxide/intrinsics/probes/div_rm_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/div_rm_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/div_rn_f32.ll b/cuda-oxide/intrinsics/probes/div_rn_f32.ll index ad3be33fc0..4c091f66c4 100644 --- a/cuda-oxide/intrinsics/probes/div_rn_f32.ll +++ b/cuda-oxide/intrinsics/probes/div_rn_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/div_rn_f64.ll b/cuda-oxide/intrinsics/probes/div_rn_f64.ll index 7716cecb0d..46cc939a94 100644 --- a/cuda-oxide/intrinsics/probes/div_rn_f64.ll +++ b/cuda-oxide/intrinsics/probes/div_rn_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/div_rn_ftz_f32.ll b/cuda-oxide/intrinsics/probes/div_rn_ftz_f32.ll index d909432081..59ead1fbf6 100644 --- a/cuda-oxide/intrinsics/probes/div_rn_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/div_rn_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/div_rp_f32.ll b/cuda-oxide/intrinsics/probes/div_rp_f32.ll index 6015dd54e1..f893b8a035 100644 --- a/cuda-oxide/intrinsics/probes/div_rp_f32.ll +++ b/cuda-oxide/intrinsics/probes/div_rp_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/div_rp_f64.ll b/cuda-oxide/intrinsics/probes/div_rp_f64.ll index d1ab996f68..8023349dd0 100644 --- a/cuda-oxide/intrinsics/probes/div_rp_f64.ll +++ b/cuda-oxide/intrinsics/probes/div_rp_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/div_rp_ftz_f32.ll b/cuda-oxide/intrinsics/probes/div_rp_ftz_f32.ll index 6457be650b..8fd724c66b 100644 --- a/cuda-oxide/intrinsics/probes/div_rp_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/div_rp_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/div_rz_f32.ll b/cuda-oxide/intrinsics/probes/div_rz_f32.ll index 3d6b30f55c..5f2d9f4ed3 100644 --- a/cuda-oxide/intrinsics/probes/div_rz_f32.ll +++ b/cuda-oxide/intrinsics/probes/div_rz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/div_rz_f64.ll b/cuda-oxide/intrinsics/probes/div_rz_f64.ll index 3681200403..a8fed6b9e3 100644 --- a/cuda-oxide/intrinsics/probes/div_rz_f64.ll +++ b/cuda-oxide/intrinsics/probes/div_rz_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/div_rz_ftz_f32.ll b/cuda-oxide/intrinsics/probes/div_rz_ftz_f32.ll index f115b4a0ac..8461f7a8cb 100644 --- a/cuda-oxide/intrinsics/probes/div_rz_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/div_rz_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/dp2a_s32.ll b/cuda-oxide/intrinsics/probes/dp2a_s32.ll index 1b520107ac..6ae168adcd 100644 --- a/cuda-oxide/intrinsics/probes/dp2a_s32.ll +++ b/cuda-oxide/intrinsics/probes/dp2a_s32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/dp2a_u32.ll b/cuda-oxide/intrinsics/probes/dp2a_u32.ll index 17b2f2100d..1a59031055 100644 --- a/cuda-oxide/intrinsics/probes/dp2a_u32.ll +++ b/cuda-oxide/intrinsics/probes/dp2a_u32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/dp4a_s32.ll b/cuda-oxide/intrinsics/probes/dp4a_s32.ll index 62741993bb..df7249d0ad 100644 --- a/cuda-oxide/intrinsics/probes/dp4a_s32.ll +++ b/cuda-oxide/intrinsics/probes/dp4a_s32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/dp4a_u32.ll b/cuda-oxide/intrinsics/probes/dp4a_u32.ll index 6977159034..11a514d7fe 100644 --- a/cuda-oxide/intrinsics/probes/dp4a_u32.ll +++ b/cuda-oxide/intrinsics/probes/dp4a_u32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/dsmem_read_u32.ll b/cuda-oxide/intrinsics/probes/dsmem_read_u32.ll index a13457b865..d3369a93c0 100644 --- a/cuda-oxide/intrinsics/probes/dsmem_read_u32.ll +++ b/cuda-oxide/intrinsics/probes/dsmem_read_u32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/dynamic_smem_size.lib_nvvm.ll b/cuda-oxide/intrinsics/probes/dynamic_smem_size.lib_nvvm.ll index 24be2c1fb8..a5cb370d22 100644 --- a/cuda-oxide/intrinsics/probes/dynamic_smem_size.lib_nvvm.ll +++ b/cuda-oxide/intrinsics/probes/dynamic_smem_size.lib_nvvm.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/dynamic_smem_size.llvm_nvptx.ll b/cuda-oxide/intrinsics/probes/dynamic_smem_size.llvm_nvptx.ll index fbda14237f..a4d63d4032 100644 --- a/cuda-oxide/intrinsics/probes/dynamic_smem_size.llvm_nvptx.ll +++ b/cuda-oxide/intrinsics/probes/dynamic_smem_size.llvm_nvptx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/elect_sync.lib_nvvm.ll b/cuda-oxide/intrinsics/probes/elect_sync.lib_nvvm.ll index b75b778ff4..b80478543a 100644 --- a/cuda-oxide/intrinsics/probes/elect_sync.lib_nvvm.ll +++ b/cuda-oxide/intrinsics/probes/elect_sync.lib_nvvm.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/elect_sync.llvm_nvptx.ll b/cuda-oxide/intrinsics/probes/elect_sync.llvm_nvptx.ll index fa7e3251e8..735bd09dbf 100644 --- a/cuda-oxide/intrinsics/probes/elect_sync.llvm_nvptx.ll +++ b/cuda-oxide/intrinsics/probes/elect_sync.llvm_nvptx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/envreg1.lib_nvvm.ll b/cuda-oxide/intrinsics/probes/envreg1.lib_nvvm.ll index 024de958ab..d00f10796b 100644 --- a/cuda-oxide/intrinsics/probes/envreg1.lib_nvvm.ll +++ b/cuda-oxide/intrinsics/probes/envreg1.lib_nvvm.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/envreg1.llvm_nvptx.ll b/cuda-oxide/intrinsics/probes/envreg1.llvm_nvptx.ll index 3b92e502af..214582c73d 100644 --- a/cuda-oxide/intrinsics/probes/envreg1.llvm_nvptx.ll +++ b/cuda-oxide/intrinsics/probes/envreg1.llvm_nvptx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/envreg2.lib_nvvm.ll b/cuda-oxide/intrinsics/probes/envreg2.lib_nvvm.ll index 3c655487c2..b9c7c0705e 100644 --- a/cuda-oxide/intrinsics/probes/envreg2.lib_nvvm.ll +++ b/cuda-oxide/intrinsics/probes/envreg2.lib_nvvm.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/envreg2.llvm_nvptx.ll b/cuda-oxide/intrinsics/probes/envreg2.llvm_nvptx.ll index c0c487c0af..13c28be699 100644 --- a/cuda-oxide/intrinsics/probes/envreg2.llvm_nvptx.ll +++ b/cuda-oxide/intrinsics/probes/envreg2.llvm_nvptx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ex2_approx_f16.ll b/cuda-oxide/intrinsics/probes/ex2_approx_f16.ll index b1539b5e8e..5817fa5d45 100644 --- a/cuda-oxide/intrinsics/probes/ex2_approx_f16.ll +++ b/cuda-oxide/intrinsics/probes/ex2_approx_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ex2_approx_f32.ll b/cuda-oxide/intrinsics/probes/ex2_approx_f32.ll index f327903a5a..3f90da483d 100644 --- a/cuda-oxide/intrinsics/probes/ex2_approx_f32.ll +++ b/cuda-oxide/intrinsics/probes/ex2_approx_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ex2_approx_ftz_f32.ll b/cuda-oxide/intrinsics/probes/ex2_approx_ftz_f32.ll index 5fa6b02a29..855c3a0bb4 100644 --- a/cuda-oxide/intrinsics/probes/ex2_approx_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/ex2_approx_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fence_acquire_shared_cluster_cluster.ll b/cuda-oxide/intrinsics/probes/fence_acquire_shared_cluster_cluster.ll new file mode 100644 index 0000000000..f4e108f093 --- /dev/null +++ b/cuda-oxide/intrinsics/probes/fence_acquire_shared_cluster_cluster.ll @@ -0,0 +1,14 @@ +; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +; SPDX-License-Identifier: Apache-2.0 +; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. +; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. + +target triple = "nvptx64-nvidia-cuda" + +define void @probe_fence_acquire_shared_cluster_cluster() #0 { + call void asm sideeffect "fence.acquire.sync_restrict::shared::cluster.cluster;", "~{memory}"() #0 + ret void +} + +attributes #0 = { convergent } diff --git a/cuda-oxide/intrinsics/probes/fence_mbarrier_init_release_cluster.ll b/cuda-oxide/intrinsics/probes/fence_mbarrier_init_release_cluster.ll index b4fdf3f7a5..8ecf0fce45 100644 --- a/cuda-oxide/intrinsics/probes/fence_mbarrier_init_release_cluster.ll +++ b/cuda-oxide/intrinsics/probes/fence_mbarrier_init_release_cluster.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fence_proxy_async_generic_acquire_shared_cluster_cluster.ll b/cuda-oxide/intrinsics/probes/fence_proxy_async_generic_acquire_shared_cluster_cluster.ll index fa972ae2a1..fdbb9bfd8b 100644 --- a/cuda-oxide/intrinsics/probes/fence_proxy_async_generic_acquire_shared_cluster_cluster.ll +++ b/cuda-oxide/intrinsics/probes/fence_proxy_async_generic_acquire_shared_cluster_cluster.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fence_proxy_async_generic_release_shared_cta_cluster.ll b/cuda-oxide/intrinsics/probes/fence_proxy_async_generic_release_shared_cta_cluster.ll index 8ca5bf56e2..8064c9814e 100644 --- a/cuda-oxide/intrinsics/probes/fence_proxy_async_generic_release_shared_cta_cluster.ll +++ b/cuda-oxide/intrinsics/probes/fence_proxy_async_generic_release_shared_cta_cluster.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fence_proxy_async_shared_cta.ll b/cuda-oxide/intrinsics/probes/fence_proxy_async_shared_cta.ll index 4f05cf5522..d3ea760911 100644 --- a/cuda-oxide/intrinsics/probes/fence_proxy_async_shared_cta.ll +++ b/cuda-oxide/intrinsics/probes/fence_proxy_async_shared_cta.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_cluster.ll b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_cluster.ll index 1431164c61..df5abfbbfb 100644 --- a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_cluster.ll +++ b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_cluster.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_cta.ll b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_cta.ll index 13cfee32c9..78f19afa06 100644 --- a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_cta.ll +++ b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_cta.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_gpu.ll b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_gpu.ll index 66df18e266..d76429b4f7 100644 --- a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_gpu.ll +++ b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_gpu.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_system.ll b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_system.ll index d936ebd25d..b851799e73 100644 --- a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_system.ll +++ b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_acquire_system.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_cluster.ll b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_cluster.ll index e61d4932d4..b53b6ca5ae 100644 --- a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_cluster.ll +++ b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_cluster.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_cta.ll b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_cta.ll index 192151092f..43603b539b 100644 --- a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_cta.ll +++ b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_cta.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_gpu.ll b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_gpu.ll index 3db5734570..60fc8a7ceb 100644 --- a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_gpu.ll +++ b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_gpu.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_system.ll b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_system.ll index d6cf138b1d..2b0f50a249 100644 --- a/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_system.ll +++ b/cuda-oxide/intrinsics/probes/fence_proxy_tensormap_generic_release_system.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_bf16x2.ll b/cuda-oxide/intrinsics/probes/fma_bf16x2.ll index 1dc117075f..576f98c0d9 100644 --- a/cuda-oxide/intrinsics/probes/fma_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/fma_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_f16x2.ll b/cuda-oxide/intrinsics/probes/fma_f16x2.ll index 54be4ac154..aee2404495 100644 --- a/cuda-oxide/intrinsics/probes/fma_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/fma_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_f32x2.ll b/cuda-oxide/intrinsics/probes/fma_f32x2.ll index 42cd1e2f34..f8452ec9a3 100644 --- a/cuda-oxide/intrinsics/probes/fma_f32x2.ll +++ b/cuda-oxide/intrinsics/probes/fma_f32x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_ftz_f16x2.ll b/cuda-oxide/intrinsics/probes/fma_ftz_f16x2.ll index 858247a169..0fcccf78f4 100644 --- a/cuda-oxide/intrinsics/probes/fma_ftz_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/fma_ftz_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_ftz_f32x2.ll b/cuda-oxide/intrinsics/probes/fma_ftz_f32x2.ll index 11f64455d3..4ba44b6905 100644 --- a/cuda-oxide/intrinsics/probes/fma_ftz_f32x2.ll +++ b/cuda-oxide/intrinsics/probes/fma_ftz_f32x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_ftz_relu_f16x2.ll b/cuda-oxide/intrinsics/probes/fma_ftz_relu_f16x2.ll index 781b2fcd8c..c74d5b7e7a 100644 --- a/cuda-oxide/intrinsics/probes/fma_ftz_relu_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/fma_ftz_relu_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_ftz_sat_f16x2.ll b/cuda-oxide/intrinsics/probes/fma_ftz_sat_f16x2.ll index 7aafe71a31..a0243d27b4 100644 --- a/cuda-oxide/intrinsics/probes/fma_ftz_sat_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/fma_ftz_sat_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_relu_bf16x2.ll b/cuda-oxide/intrinsics/probes/fma_relu_bf16x2.ll index 527cae7873..0374c45d22 100644 --- a/cuda-oxide/intrinsics/probes/fma_relu_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/fma_relu_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_relu_f16x2.ll b/cuda-oxide/intrinsics/probes/fma_relu_f16x2.ll index 01aed4b881..dc6d1a4754 100644 --- a/cuda-oxide/intrinsics/probes/fma_relu_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/fma_relu_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rm_f32.ll b/cuda-oxide/intrinsics/probes/fma_rm_f32.ll index d2341d092b..12a7ce41aa 100644 --- a/cuda-oxide/intrinsics/probes/fma_rm_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rm_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rm_f64.ll b/cuda-oxide/intrinsics/probes/fma_rm_f64.ll index 6fd0105b2a..1cd95c9542 100644 --- a/cuda-oxide/intrinsics/probes/fma_rm_f64.ll +++ b/cuda-oxide/intrinsics/probes/fma_rm_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rm_ftz_f32.ll b/cuda-oxide/intrinsics/probes/fma_rm_ftz_f32.ll index 2573e5d230..8d94fe60a8 100644 --- a/cuda-oxide/intrinsics/probes/fma_rm_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rm_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rm_ftz_sat_f32.ll b/cuda-oxide/intrinsics/probes/fma_rm_ftz_sat_f32.ll index 2ed0f51884..1e7917ba74 100644 --- a/cuda-oxide/intrinsics/probes/fma_rm_ftz_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rm_ftz_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rm_sat_f32.ll b/cuda-oxide/intrinsics/probes/fma_rm_sat_f32.ll index 76decf7520..711811dac5 100644 --- a/cuda-oxide/intrinsics/probes/fma_rm_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rm_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rn_f32.ll b/cuda-oxide/intrinsics/probes/fma_rn_f32.ll index 942623c239..82ae50a12f 100644 --- a/cuda-oxide/intrinsics/probes/fma_rn_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rn_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rn_f64.ll b/cuda-oxide/intrinsics/probes/fma_rn_f64.ll index 7edeea3bde..5952e2e2b3 100644 --- a/cuda-oxide/intrinsics/probes/fma_rn_f64.ll +++ b/cuda-oxide/intrinsics/probes/fma_rn_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rn_ftz_f32.ll b/cuda-oxide/intrinsics/probes/fma_rn_ftz_f32.ll index 9728b33585..289cfb935b 100644 --- a/cuda-oxide/intrinsics/probes/fma_rn_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rn_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rn_ftz_sat_f32.ll b/cuda-oxide/intrinsics/probes/fma_rn_ftz_sat_f32.ll index 68f0cf2fdc..a96e72c2fd 100644 --- a/cuda-oxide/intrinsics/probes/fma_rn_ftz_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rn_ftz_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rn_sat_f32.ll b/cuda-oxide/intrinsics/probes/fma_rn_sat_f32.ll index f431a42449..c35417ff00 100644 --- a/cuda-oxide/intrinsics/probes/fma_rn_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rn_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rp_f32.ll b/cuda-oxide/intrinsics/probes/fma_rp_f32.ll index 32003d3aef..76f9ff3c5c 100644 --- a/cuda-oxide/intrinsics/probes/fma_rp_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rp_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rp_f64.ll b/cuda-oxide/intrinsics/probes/fma_rp_f64.ll index 304eaf0781..7a92fd5824 100644 --- a/cuda-oxide/intrinsics/probes/fma_rp_f64.ll +++ b/cuda-oxide/intrinsics/probes/fma_rp_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rp_ftz_f32.ll b/cuda-oxide/intrinsics/probes/fma_rp_ftz_f32.ll index bbb47a8e1a..96f89cb7b9 100644 --- a/cuda-oxide/intrinsics/probes/fma_rp_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rp_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rp_ftz_sat_f32.ll b/cuda-oxide/intrinsics/probes/fma_rp_ftz_sat_f32.ll index 6cca7ee08b..6922783d07 100644 --- a/cuda-oxide/intrinsics/probes/fma_rp_ftz_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rp_ftz_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rp_sat_f32.ll b/cuda-oxide/intrinsics/probes/fma_rp_sat_f32.ll index ce9ac06555..70d66d7ca9 100644 --- a/cuda-oxide/intrinsics/probes/fma_rp_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rp_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rz_f32.ll b/cuda-oxide/intrinsics/probes/fma_rz_f32.ll index b1c2cf0bdf..9bc41214da 100644 --- a/cuda-oxide/intrinsics/probes/fma_rz_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rz_f64.ll b/cuda-oxide/intrinsics/probes/fma_rz_f64.ll index 62c022fe48..8198cc63a2 100644 --- a/cuda-oxide/intrinsics/probes/fma_rz_f64.ll +++ b/cuda-oxide/intrinsics/probes/fma_rz_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rz_ftz_f32.ll b/cuda-oxide/intrinsics/probes/fma_rz_ftz_f32.ll index aa2c5ce792..21599e9afc 100644 --- a/cuda-oxide/intrinsics/probes/fma_rz_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rz_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rz_ftz_sat_f32.ll b/cuda-oxide/intrinsics/probes/fma_rz_ftz_sat_f32.ll index bd490ce9a0..df4b05ab00 100644 --- a/cuda-oxide/intrinsics/probes/fma_rz_ftz_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rz_ftz_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_rz_sat_f32.ll b/cuda-oxide/intrinsics/probes/fma_rz_sat_f32.ll index 079951980f..30bac3cb8c 100644 --- a/cuda-oxide/intrinsics/probes/fma_rz_sat_f32.ll +++ b/cuda-oxide/intrinsics/probes/fma_rz_sat_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/fma_sat_f16x2.ll b/cuda-oxide/intrinsics/probes/fma_sat_f16x2.ll index e42116c026..ab31a8df71 100644 --- a/cuda-oxide/intrinsics/probes/fma_sat_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/fma_sat_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/globaltimer.lib_nvvm.ll b/cuda-oxide/intrinsics/probes/globaltimer.lib_nvvm.ll index e42da9d97d..67ce88aaea 100644 --- a/cuda-oxide/intrinsics/probes/globaltimer.lib_nvvm.ll +++ b/cuda-oxide/intrinsics/probes/globaltimer.lib_nvvm.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/globaltimer.llvm_nvptx.ll b/cuda-oxide/intrinsics/probes/globaltimer.llvm_nvptx.ll index 90b07fb3e6..027957c8b8 100644 --- a/cuda-oxide/intrinsics/probes/globaltimer.llvm_nvptx.ll +++ b/cuda-oxide/intrinsics/probes/globaltimer.llvm_nvptx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/grid_dependency_launch_dependents.ll b/cuda-oxide/intrinsics/probes/grid_dependency_launch_dependents.ll index 32cd3ccaa9..54ec00a8ae 100644 --- a/cuda-oxide/intrinsics/probes/grid_dependency_launch_dependents.ll +++ b/cuda-oxide/intrinsics/probes/grid_dependency_launch_dependents.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/grid_dependency_wait.ll b/cuda-oxide/intrinsics/probes/grid_dependency_wait.ll index ff6e1d4556..777d9e47b7 100644 --- a/cuda-oxide/intrinsics/probes/grid_dependency_wait.ll +++ b/cuda-oxide/intrinsics/probes/grid_dependency_wait.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/grid_dim_x.ll b/cuda-oxide/intrinsics/probes/grid_dim_x.ll index 90d1ff3b41..24c0838a8f 100644 --- a/cuda-oxide/intrinsics/probes/grid_dim_x.ll +++ b/cuda-oxide/intrinsics/probes/grid_dim_x.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/grid_dim_y.ll b/cuda-oxide/intrinsics/probes/grid_dim_y.ll index 834ebe97e2..0bfe83ec60 100644 --- a/cuda-oxide/intrinsics/probes/grid_dim_y.ll +++ b/cuda-oxide/intrinsics/probes/grid_dim_y.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/grid_dim_z.ll b/cuda-oxide/intrinsics/probes/grid_dim_z.ll index 9bdb359a11..ffa379d7ee 100644 --- a/cuda-oxide/intrinsics/probes/grid_dim_z.ll +++ b/cuda-oxide/intrinsics/probes/grid_dim_z.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/gridid.lib_nvvm.ll b/cuda-oxide/intrinsics/probes/gridid.lib_nvvm.ll index c775175f32..d9e008506f 100644 --- a/cuda-oxide/intrinsics/probes/gridid.lib_nvvm.ll +++ b/cuda-oxide/intrinsics/probes/gridid.lib_nvvm.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/gridid.llvm_nvptx.ll b/cuda-oxide/intrinsics/probes/gridid.llvm_nvptx.ll index 9cc9d2b3e1..9d356d339e 100644 --- a/cuda-oxide/intrinsics/probes/gridid.llvm_nvptx.ll +++ b/cuda-oxide/intrinsics/probes/gridid.llvm_nvptx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/lane_id.ll b/cuda-oxide/intrinsics/probes/lane_id.ll index 7bbac3d41e..9474603dec 100644 --- a/cuda-oxide/intrinsics/probes/lane_id.ll +++ b/cuda-oxide/intrinsics/probes/lane_id.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/lanemask_eq.ll b/cuda-oxide/intrinsics/probes/lanemask_eq.ll index e9e49ebcaf..df365791ac 100644 --- a/cuda-oxide/intrinsics/probes/lanemask_eq.ll +++ b/cuda-oxide/intrinsics/probes/lanemask_eq.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/lanemask_ge.ll b/cuda-oxide/intrinsics/probes/lanemask_ge.ll index 4a7af9e71b..d62495bf6e 100644 --- a/cuda-oxide/intrinsics/probes/lanemask_ge.ll +++ b/cuda-oxide/intrinsics/probes/lanemask_ge.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/lanemask_gt.ll b/cuda-oxide/intrinsics/probes/lanemask_gt.ll index 46fca59a6d..bbd63ad1b5 100644 --- a/cuda-oxide/intrinsics/probes/lanemask_gt.ll +++ b/cuda-oxide/intrinsics/probes/lanemask_gt.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/lanemask_le.ll b/cuda-oxide/intrinsics/probes/lanemask_le.ll index 69702617cc..f91d4e2240 100644 --- a/cuda-oxide/intrinsics/probes/lanemask_le.ll +++ b/cuda-oxide/intrinsics/probes/lanemask_le.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/lanemask_lt.ll b/cuda-oxide/intrinsics/probes/lanemask_lt.ll index 9d7e9df70d..71ddf8630d 100644 --- a/cuda-oxide/intrinsics/probes/lanemask_lt.ll +++ b/cuda-oxide/intrinsics/probes/lanemask_lt.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x1_trans_b8.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x1_trans_b8.ll index d7120560bd..a107af06ff 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x1_trans_b8.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x1_trans_b8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x1_trans_b8x16_b4x16_p64.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x1_trans_b8x16_b4x16_p64.ll index 71d4194647..2f9996f926 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x1_trans_b8x16_b4x16_p64.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x1_trans_b8x16_b4x16_p64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x1_trans_b8x16_b6x16_p32.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x1_trans_b8x16_b6x16_p32.ll index 63d721fba8..17feddc176 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x1_trans_b8x16_b6x16_p32.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x1_trans_b8x16_b6x16_p32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x2_trans_b8.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x2_trans_b8.ll index 03e3b56096..cedb8d3f97 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x2_trans_b8.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x2_trans_b8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x2_trans_b8x16_b4x16_p64.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x2_trans_b8x16_b4x16_p64.ll index 12f331dc2f..fbb75ea3b8 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x2_trans_b8x16_b4x16_p64.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x2_trans_b8x16_b4x16_p64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x2_trans_b8x16_b6x16_p32.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x2_trans_b8x16_b6x16_p32.ll index fffc9ee48d..f05fe7c01d 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x2_trans_b8x16_b6x16_p32.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m16n16_x2_trans_b8x16_b6x16_p32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x1_b8x16_b4x16_p64.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x1_b8x16_b4x16_p64.ll index 08404013e1..032b1e49ef 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x1_b8x16_b4x16_p64.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x1_b8x16_b4x16_p64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x1_b8x16_b6x16_p32.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x1_b8x16_b6x16_p32.ll index e26ad58d9a..71737464a1 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x1_b8x16_b6x16_p32.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x1_b8x16_b6x16_p32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x2_b8x16_b4x16_p64.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x2_b8x16_b4x16_p64.ll index 44b877cb36..d405a655a0 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x2_b8x16_b4x16_p64.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x2_b8x16_b4x16_p64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x2_b8x16_b6x16_p32.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x2_b8x16_b6x16_p32.ll index f99583078c..0c91d424a5 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x2_b8x16_b6x16_p32.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x2_b8x16_b6x16_p32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x4_b8x16_b4x16_p64.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x4_b8x16_b4x16_p64.ll index 76db3c68a5..6a82c440d6 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x4_b8x16_b4x16_p64.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x4_b8x16_b4x16_p64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x4_b8x16_b6x16_p32.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x4_b8x16_b6x16_p32.ll index a97c39e959..7b4befab76 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x4_b8x16_b6x16_p32.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m8n16_x4_b8x16_b6x16_p32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x1_b16.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x1_b16.ll index 0df161d491..4ad20b1489 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x1_b16.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x1_b16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x1_trans_b16.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x1_trans_b16.ll index 4aa419f294..7435d42857 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x1_trans_b16.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x1_trans_b16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x2_b16.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x2_b16.ll index 620c9e8503..122e6337f5 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x2_b16.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x2_b16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x2_trans_b16.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x2_trans_b16.ll index 276b306924..2f4b0fb19a 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x2_trans_b16.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x2_trans_b16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x4_b16.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x4_b16.ll index 8dee2a1596..9afdefb05e 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x4_b16.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x4_b16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x4_trans_b16.ll b/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x4_trans_b16.ll index ca9bb5beff..ed833242bc 100644 --- a/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x4_trans_b16.ll +++ b/cuda-oxide/intrinsics/probes/ldmatrix_m8n8_x4_trans_b16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/lg2_approx_f32.ll b/cuda-oxide/intrinsics/probes/lg2_approx_f32.ll index 4e9ae76a2b..a241b9a103 100644 --- a/cuda-oxide/intrinsics/probes/lg2_approx_f32.ll +++ b/cuda-oxide/intrinsics/probes/lg2_approx_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/lg2_approx_ftz_f32.ll b/cuda-oxide/intrinsics/probes/lg2_approx_ftz_f32.ll index 1f72f7fb57..955b4da06c 100644 --- a/cuda-oxide/intrinsics/probes/lg2_approx_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/lg2_approx_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/map_shared_rank.ll b/cuda-oxide/intrinsics/probes/map_shared_rank.ll index 97c4b494d0..0c9c7e34d5 100644 --- a/cuda-oxide/intrinsics/probes/map_shared_rank.ll +++ b/cuda-oxide/intrinsics/probes/map_shared_rank.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/match_all_i64_sync.ll b/cuda-oxide/intrinsics/probes/match_all_i64_sync.ll index 1037b36dde..85af56e153 100644 --- a/cuda-oxide/intrinsics/probes/match_all_i64_sync.ll +++ b/cuda-oxide/intrinsics/probes/match_all_i64_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/match_all_sync.ll b/cuda-oxide/intrinsics/probes/match_all_sync.ll index a1e42a5025..f4524cf3d7 100644 --- a/cuda-oxide/intrinsics/probes/match_all_sync.ll +++ b/cuda-oxide/intrinsics/probes/match_all_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/match_any_i64_sync.ll b/cuda-oxide/intrinsics/probes/match_any_i64_sync.ll index 3adeb7e71c..794ef63ee4 100644 --- a/cuda-oxide/intrinsics/probes/match_any_i64_sync.ll +++ b/cuda-oxide/intrinsics/probes/match_any_i64_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/match_any_sync.ll b/cuda-oxide/intrinsics/probes/match_any_sync.ll index c444c5d0c5..c99b3b1faf 100644 --- a/cuda-oxide/intrinsics/probes/match_any_sync.ll +++ b/cuda-oxide/intrinsics/probes/match_any_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_bf16.ll b/cuda-oxide/intrinsics/probes/max_bf16.ll index a7f9ecf213..c08f0aba8f 100644 --- a/cuda-oxide/intrinsics/probes/max_bf16.ll +++ b/cuda-oxide/intrinsics/probes/max_bf16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_bf16x2.ll b/cuda-oxide/intrinsics/probes/max_bf16x2.ll index 1cfe1a29dd..f1891d5f7b 100644 --- a/cuda-oxide/intrinsics/probes/max_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_f16.ll b/cuda-oxide/intrinsics/probes/max_f16.ll index 13ee207f89..56b8d32df3 100644 --- a/cuda-oxide/intrinsics/probes/max_f16.ll +++ b/cuda-oxide/intrinsics/probes/max_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_f16x2.ll b/cuda-oxide/intrinsics/probes/max_f16x2.ll index 5cb09e75f9..e23de67806 100644 --- a/cuda-oxide/intrinsics/probes/max_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_ftz_f16.ll b/cuda-oxide/intrinsics/probes/max_ftz_f16.ll index 6968a87710..07ab40ca8e 100644 --- a/cuda-oxide/intrinsics/probes/max_ftz_f16.ll +++ b/cuda-oxide/intrinsics/probes/max_ftz_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_ftz_f16x2.ll b/cuda-oxide/intrinsics/probes/max_ftz_f16x2.ll index 35ebd53300..51d2683bc3 100644 --- a/cuda-oxide/intrinsics/probes/max_ftz_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_ftz_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_ftz_nan_f16.ll b/cuda-oxide/intrinsics/probes/max_ftz_nan_f16.ll index 62bd90e757..5c01e99e74 100644 --- a/cuda-oxide/intrinsics/probes/max_ftz_nan_f16.ll +++ b/cuda-oxide/intrinsics/probes/max_ftz_nan_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_ftz_nan_f16x2.ll b/cuda-oxide/intrinsics/probes/max_ftz_nan_f16x2.ll index 859faca949..8b4e2cd4ef 100644 --- a/cuda-oxide/intrinsics/probes/max_ftz_nan_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_ftz_nan_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_ftz_nan_xorsign_abs_f16.ll b/cuda-oxide/intrinsics/probes/max_ftz_nan_xorsign_abs_f16.ll index a6bcab50cb..9f25c9b9f1 100644 --- a/cuda-oxide/intrinsics/probes/max_ftz_nan_xorsign_abs_f16.ll +++ b/cuda-oxide/intrinsics/probes/max_ftz_nan_xorsign_abs_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_ftz_nan_xorsign_abs_f16x2.ll b/cuda-oxide/intrinsics/probes/max_ftz_nan_xorsign_abs_f16x2.ll index 518144a909..9e03853650 100644 --- a/cuda-oxide/intrinsics/probes/max_ftz_nan_xorsign_abs_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_ftz_nan_xorsign_abs_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_ftz_nan_xorsign_abs_f32.ll b/cuda-oxide/intrinsics/probes/max_ftz_nan_xorsign_abs_f32.ll index 1ee93b2b28..bd52f41e68 100644 --- a/cuda-oxide/intrinsics/probes/max_ftz_nan_xorsign_abs_f32.ll +++ b/cuda-oxide/intrinsics/probes/max_ftz_nan_xorsign_abs_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_ftz_xorsign_abs_f16.ll b/cuda-oxide/intrinsics/probes/max_ftz_xorsign_abs_f16.ll index a667ecbada..80a53da299 100644 --- a/cuda-oxide/intrinsics/probes/max_ftz_xorsign_abs_f16.ll +++ b/cuda-oxide/intrinsics/probes/max_ftz_xorsign_abs_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_ftz_xorsign_abs_f16x2.ll b/cuda-oxide/intrinsics/probes/max_ftz_xorsign_abs_f16x2.ll index 8d97536f8a..ebd68a6021 100644 --- a/cuda-oxide/intrinsics/probes/max_ftz_xorsign_abs_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_ftz_xorsign_abs_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_ftz_xorsign_abs_f32.ll b/cuda-oxide/intrinsics/probes/max_ftz_xorsign_abs_f32.ll index bf5deaecc1..480062707c 100644 --- a/cuda-oxide/intrinsics/probes/max_ftz_xorsign_abs_f32.ll +++ b/cuda-oxide/intrinsics/probes/max_ftz_xorsign_abs_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_nan_bf16.ll b/cuda-oxide/intrinsics/probes/max_nan_bf16.ll index 46951929db..4b66fe262b 100644 --- a/cuda-oxide/intrinsics/probes/max_nan_bf16.ll +++ b/cuda-oxide/intrinsics/probes/max_nan_bf16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_nan_bf16x2.ll b/cuda-oxide/intrinsics/probes/max_nan_bf16x2.ll index 6d4b0656dd..c4cd04ccb8 100644 --- a/cuda-oxide/intrinsics/probes/max_nan_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_nan_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_nan_f16.ll b/cuda-oxide/intrinsics/probes/max_nan_f16.ll index e80fc55c42..323316af0f 100644 --- a/cuda-oxide/intrinsics/probes/max_nan_f16.ll +++ b/cuda-oxide/intrinsics/probes/max_nan_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_nan_f16x2.ll b/cuda-oxide/intrinsics/probes/max_nan_f16x2.ll index fffb412bcf..3b59c76fd8 100644 --- a/cuda-oxide/intrinsics/probes/max_nan_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_nan_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_bf16.ll b/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_bf16.ll index 46520aa757..b56db8883c 100644 --- a/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_bf16.ll +++ b/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_bf16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_bf16x2.ll b/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_bf16x2.ll index c2448f9493..d70370f02e 100644 --- a/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_f16.ll b/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_f16.ll index f5a4f73863..8360b06d31 100644 --- a/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_f16.ll +++ b/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_f16x2.ll b/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_f16x2.ll index e83ff6d073..c62c607866 100644 --- a/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_f32.ll b/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_f32.ll index 1d5d4c2d05..583c2601ee 100644 --- a/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_f32.ll +++ b/cuda-oxide/intrinsics/probes/max_nan_xorsign_abs_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_relu_s16x2.ll b/cuda-oxide/intrinsics/probes/max_relu_s16x2.ll index ac74b1eb86..72135affff 100644 --- a/cuda-oxide/intrinsics/probes/max_relu_s16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_relu_s16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_relu_s32.ll b/cuda-oxide/intrinsics/probes/max_relu_s32.ll index 2899a43e88..6e276ae7f9 100644 --- a/cuda-oxide/intrinsics/probes/max_relu_s32.ll +++ b/cuda-oxide/intrinsics/probes/max_relu_s32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_s16x2.ll b/cuda-oxide/intrinsics/probes/max_s16x2.ll index 7cfbe1d6ec..843b7e6359 100644 --- a/cuda-oxide/intrinsics/probes/max_s16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_s16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_u16x2.ll b/cuda-oxide/intrinsics/probes/max_u16x2.ll index 81709297a8..91fcfb2bfd 100644 --- a/cuda-oxide/intrinsics/probes/max_u16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_u16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_xorsign_abs_bf16.ll b/cuda-oxide/intrinsics/probes/max_xorsign_abs_bf16.ll index b953a30809..0d657e8a08 100644 --- a/cuda-oxide/intrinsics/probes/max_xorsign_abs_bf16.ll +++ b/cuda-oxide/intrinsics/probes/max_xorsign_abs_bf16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_xorsign_abs_bf16x2.ll b/cuda-oxide/intrinsics/probes/max_xorsign_abs_bf16x2.ll index af3d3cd3e8..1ce8ab3e75 100644 --- a/cuda-oxide/intrinsics/probes/max_xorsign_abs_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_xorsign_abs_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_xorsign_abs_f16.ll b/cuda-oxide/intrinsics/probes/max_xorsign_abs_f16.ll index e57d67dce1..76340b8a69 100644 --- a/cuda-oxide/intrinsics/probes/max_xorsign_abs_f16.ll +++ b/cuda-oxide/intrinsics/probes/max_xorsign_abs_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_xorsign_abs_f16x2.ll b/cuda-oxide/intrinsics/probes/max_xorsign_abs_f16x2.ll index df4064baa0..a1e63430f2 100644 --- a/cuda-oxide/intrinsics/probes/max_xorsign_abs_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/max_xorsign_abs_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/max_xorsign_abs_f32.ll b/cuda-oxide/intrinsics/probes/max_xorsign_abs_f32.ll index 88f55ad94b..4000f38a09 100644 --- a/cuda-oxide/intrinsics/probes/max_xorsign_abs_f32.ll +++ b/cuda-oxide/intrinsics/probes/max_xorsign_abs_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mbarrier_arrive.ll b/cuda-oxide/intrinsics/probes/mbarrier_arrive.ll index 00887dc0e1..1936149f29 100644 --- a/cuda-oxide/intrinsics/probes/mbarrier_arrive.ll +++ b/cuda-oxide/intrinsics/probes/mbarrier_arrive.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mbarrier_arrive_cluster.ll b/cuda-oxide/intrinsics/probes/mbarrier_arrive_cluster.ll index 358d5f397e..d1af2d70ab 100644 --- a/cuda-oxide/intrinsics/probes/mbarrier_arrive_cluster.ll +++ b/cuda-oxide/intrinsics/probes/mbarrier_arrive_cluster.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mbarrier_arrive_cluster_relaxed.ll b/cuda-oxide/intrinsics/probes/mbarrier_arrive_cluster_relaxed.ll new file mode 100644 index 0000000000..7803407040 --- /dev/null +++ b/cuda-oxide/intrinsics/probes/mbarrier_arrive_cluster_relaxed.ll @@ -0,0 +1,14 @@ +; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +; SPDX-License-Identifier: Apache-2.0 +; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. +; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. + +target triple = "nvptx64-nvidia-cuda" + +define void @probe_mbarrier_arrive_cluster_relaxed(i64 %address) #0 { + call void asm sideeffect "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [$0];", "l,~{memory}"(i64 %address) #0 + ret void +} + +attributes #0 = { convergent } diff --git a/cuda-oxide/intrinsics/probes/mbarrier_arrive_expect_tx.ll b/cuda-oxide/intrinsics/probes/mbarrier_arrive_expect_tx.ll index 5ac7f29c45..09cdb4c39b 100644 --- a/cuda-oxide/intrinsics/probes/mbarrier_arrive_expect_tx.ll +++ b/cuda-oxide/intrinsics/probes/mbarrier_arrive_expect_tx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mbarrier_arrive_expect_tx_cluster.ll b/cuda-oxide/intrinsics/probes/mbarrier_arrive_expect_tx_cluster.ll index 21b385104c..36f0399722 100644 --- a/cuda-oxide/intrinsics/probes/mbarrier_arrive_expect_tx_cluster.ll +++ b/cuda-oxide/intrinsics/probes/mbarrier_arrive_expect_tx_cluster.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mbarrier_arrive_no_complete.ll b/cuda-oxide/intrinsics/probes/mbarrier_arrive_no_complete.ll index c71b13f953..e12561f808 100644 --- a/cuda-oxide/intrinsics/probes/mbarrier_arrive_no_complete.ll +++ b/cuda-oxide/intrinsics/probes/mbarrier_arrive_no_complete.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mbarrier_init.ll b/cuda-oxide/intrinsics/probes/mbarrier_init.ll index 83bad63b9b..d1ed4e781e 100644 --- a/cuda-oxide/intrinsics/probes/mbarrier_init.ll +++ b/cuda-oxide/intrinsics/probes/mbarrier_init.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mbarrier_inval.ll b/cuda-oxide/intrinsics/probes/mbarrier_inval.ll index 63a7fecee3..e4f71151dc 100644 --- a/cuda-oxide/intrinsics/probes/mbarrier_inval.ll +++ b/cuda-oxide/intrinsics/probes/mbarrier_inval.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mbarrier_test_wait.ll b/cuda-oxide/intrinsics/probes/mbarrier_test_wait.ll index 8d09e344aa..a536c93b77 100644 --- a/cuda-oxide/intrinsics/probes/mbarrier_test_wait.ll +++ b/cuda-oxide/intrinsics/probes/mbarrier_test_wait.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mbarrier_try_wait.ll b/cuda-oxide/intrinsics/probes/mbarrier_try_wait.ll index da9443d83f..aa5e81c901 100644 --- a/cuda-oxide/intrinsics/probes/mbarrier_try_wait.ll +++ b/cuda-oxide/intrinsics/probes/mbarrier_try_wait.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mbarrier_try_wait_parity.ll b/cuda-oxide/intrinsics/probes/mbarrier_try_wait_parity.ll index 7c5f94a576..3e166021f3 100644 --- a/cuda-oxide/intrinsics/probes/mbarrier_try_wait_parity.ll +++ b/cuda-oxide/intrinsics/probes/mbarrier_try_wait_parity.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mbarrier_try_wait_parity_cluster.ll b/cuda-oxide/intrinsics/probes/mbarrier_try_wait_parity_cluster.ll index bf4f0ab82a..d86ac07669 100644 --- a/cuda-oxide/intrinsics/probes/mbarrier_try_wait_parity_cluster.ll +++ b/cuda-oxide/intrinsics/probes/mbarrier_try_wait_parity_cluster.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mbarrier_try_wait_parity_cluster_relaxed.ll b/cuda-oxide/intrinsics/probes/mbarrier_try_wait_parity_cluster_relaxed.ll new file mode 100644 index 0000000000..beb9018917 --- /dev/null +++ b/cuda-oxide/intrinsics/probes/mbarrier_try_wait_parity_cluster_relaxed.ll @@ -0,0 +1,16 @@ +; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +; SPDX-License-Identifier: Apache-2.0 +; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. +; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. + +target triple = "nvptx64-nvidia-cuda" + +define i1 @probe_mbarrier_try_wait_parity_cluster_relaxed(ptr %barrier_generic, i32 %parity) #0 { + %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3) + %result_i32 = call i32 asm sideeffect "{ .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}"(ptr addrspace(3) %barrier, i32 %parity) #0 + %result = trunc i32 %result_i32 to i1 + ret i1 %result +} + +attributes #0 = { convergent } diff --git a/cuda-oxide/intrinsics/probes/min_bf16.ll b/cuda-oxide/intrinsics/probes/min_bf16.ll index b12b994fe9..1b97c96f54 100644 --- a/cuda-oxide/intrinsics/probes/min_bf16.ll +++ b/cuda-oxide/intrinsics/probes/min_bf16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_bf16x2.ll b/cuda-oxide/intrinsics/probes/min_bf16x2.ll index 755f261611..8bd5fa44bc 100644 --- a/cuda-oxide/intrinsics/probes/min_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_f16.ll b/cuda-oxide/intrinsics/probes/min_f16.ll index b81d2a97cb..6b4c461e91 100644 --- a/cuda-oxide/intrinsics/probes/min_f16.ll +++ b/cuda-oxide/intrinsics/probes/min_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_f16x2.ll b/cuda-oxide/intrinsics/probes/min_f16x2.ll index 50371fb43d..0ac0766fb9 100644 --- a/cuda-oxide/intrinsics/probes/min_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_ftz_f16.ll b/cuda-oxide/intrinsics/probes/min_ftz_f16.ll index 46a4161a13..bc2048739b 100644 --- a/cuda-oxide/intrinsics/probes/min_ftz_f16.ll +++ b/cuda-oxide/intrinsics/probes/min_ftz_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_ftz_f16x2.ll b/cuda-oxide/intrinsics/probes/min_ftz_f16x2.ll index 847d786cf0..9ebad9a53d 100644 --- a/cuda-oxide/intrinsics/probes/min_ftz_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_ftz_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_ftz_nan_f16.ll b/cuda-oxide/intrinsics/probes/min_ftz_nan_f16.ll index 71fa562e7e..a4980bae2b 100644 --- a/cuda-oxide/intrinsics/probes/min_ftz_nan_f16.ll +++ b/cuda-oxide/intrinsics/probes/min_ftz_nan_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_ftz_nan_f16x2.ll b/cuda-oxide/intrinsics/probes/min_ftz_nan_f16x2.ll index 7b6f4c03b2..b04a956423 100644 --- a/cuda-oxide/intrinsics/probes/min_ftz_nan_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_ftz_nan_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_ftz_nan_xorsign_abs_f16.ll b/cuda-oxide/intrinsics/probes/min_ftz_nan_xorsign_abs_f16.ll index f5adfc61db..13d442ef7e 100644 --- a/cuda-oxide/intrinsics/probes/min_ftz_nan_xorsign_abs_f16.ll +++ b/cuda-oxide/intrinsics/probes/min_ftz_nan_xorsign_abs_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_ftz_nan_xorsign_abs_f16x2.ll b/cuda-oxide/intrinsics/probes/min_ftz_nan_xorsign_abs_f16x2.ll index 6f571b357b..7d5ab767db 100644 --- a/cuda-oxide/intrinsics/probes/min_ftz_nan_xorsign_abs_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_ftz_nan_xorsign_abs_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_ftz_nan_xorsign_abs_f32.ll b/cuda-oxide/intrinsics/probes/min_ftz_nan_xorsign_abs_f32.ll index 6f6d29c611..39d122a009 100644 --- a/cuda-oxide/intrinsics/probes/min_ftz_nan_xorsign_abs_f32.ll +++ b/cuda-oxide/intrinsics/probes/min_ftz_nan_xorsign_abs_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_ftz_xorsign_abs_f16.ll b/cuda-oxide/intrinsics/probes/min_ftz_xorsign_abs_f16.ll index c55d1815bf..aaf1e36906 100644 --- a/cuda-oxide/intrinsics/probes/min_ftz_xorsign_abs_f16.ll +++ b/cuda-oxide/intrinsics/probes/min_ftz_xorsign_abs_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_ftz_xorsign_abs_f16x2.ll b/cuda-oxide/intrinsics/probes/min_ftz_xorsign_abs_f16x2.ll index 0ae17100ee..e216847047 100644 --- a/cuda-oxide/intrinsics/probes/min_ftz_xorsign_abs_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_ftz_xorsign_abs_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_ftz_xorsign_abs_f32.ll b/cuda-oxide/intrinsics/probes/min_ftz_xorsign_abs_f32.ll index dfb5e4574d..e319e8bd31 100644 --- a/cuda-oxide/intrinsics/probes/min_ftz_xorsign_abs_f32.ll +++ b/cuda-oxide/intrinsics/probes/min_ftz_xorsign_abs_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_nan_bf16.ll b/cuda-oxide/intrinsics/probes/min_nan_bf16.ll index a9d7cb70c8..585c5fdd63 100644 --- a/cuda-oxide/intrinsics/probes/min_nan_bf16.ll +++ b/cuda-oxide/intrinsics/probes/min_nan_bf16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_nan_bf16x2.ll b/cuda-oxide/intrinsics/probes/min_nan_bf16x2.ll index db8dc9f47c..7caf9832ee 100644 --- a/cuda-oxide/intrinsics/probes/min_nan_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_nan_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_nan_f16.ll b/cuda-oxide/intrinsics/probes/min_nan_f16.ll index 2e96b8e9ed..20406393cc 100644 --- a/cuda-oxide/intrinsics/probes/min_nan_f16.ll +++ b/cuda-oxide/intrinsics/probes/min_nan_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_nan_f16x2.ll b/cuda-oxide/intrinsics/probes/min_nan_f16x2.ll index 2ed9697f0b..5de0b7d656 100644 --- a/cuda-oxide/intrinsics/probes/min_nan_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_nan_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_bf16.ll b/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_bf16.ll index ed15873b46..3d7683a079 100644 --- a/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_bf16.ll +++ b/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_bf16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_bf16x2.ll b/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_bf16x2.ll index e4c16ad8c1..266db92e11 100644 --- a/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_f16.ll b/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_f16.ll index 163461458b..fd1f221881 100644 --- a/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_f16.ll +++ b/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_f16x2.ll b/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_f16x2.ll index 7c0cda0722..6490aef267 100644 --- a/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_f32.ll b/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_f32.ll index e5ec8cd195..11484c72d2 100644 --- a/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_f32.ll +++ b/cuda-oxide/intrinsics/probes/min_nan_xorsign_abs_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_relu_s16x2.ll b/cuda-oxide/intrinsics/probes/min_relu_s16x2.ll index 7cc66f1708..3911e06e90 100644 --- a/cuda-oxide/intrinsics/probes/min_relu_s16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_relu_s16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_relu_s32.ll b/cuda-oxide/intrinsics/probes/min_relu_s32.ll index d2731eaaba..4ade2ba0f0 100644 --- a/cuda-oxide/intrinsics/probes/min_relu_s32.ll +++ b/cuda-oxide/intrinsics/probes/min_relu_s32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_s16x2.ll b/cuda-oxide/intrinsics/probes/min_s16x2.ll index 1df798cb24..80005dc928 100644 --- a/cuda-oxide/intrinsics/probes/min_s16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_s16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_u16x2.ll b/cuda-oxide/intrinsics/probes/min_u16x2.ll index 1ccd9ec236..50577f1925 100644 --- a/cuda-oxide/intrinsics/probes/min_u16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_u16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_xorsign_abs_bf16.ll b/cuda-oxide/intrinsics/probes/min_xorsign_abs_bf16.ll index 991b37c529..4d85072533 100644 --- a/cuda-oxide/intrinsics/probes/min_xorsign_abs_bf16.ll +++ b/cuda-oxide/intrinsics/probes/min_xorsign_abs_bf16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_xorsign_abs_bf16x2.ll b/cuda-oxide/intrinsics/probes/min_xorsign_abs_bf16x2.ll index f6f4bb7960..6fb0c440ba 100644 --- a/cuda-oxide/intrinsics/probes/min_xorsign_abs_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_xorsign_abs_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_xorsign_abs_f16.ll b/cuda-oxide/intrinsics/probes/min_xorsign_abs_f16.ll index cee6a4908a..3953ade735 100644 --- a/cuda-oxide/intrinsics/probes/min_xorsign_abs_f16.ll +++ b/cuda-oxide/intrinsics/probes/min_xorsign_abs_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_xorsign_abs_f16x2.ll b/cuda-oxide/intrinsics/probes/min_xorsign_abs_f16x2.ll index 14385038c1..8004253fc6 100644 --- a/cuda-oxide/intrinsics/probes/min_xorsign_abs_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/min_xorsign_abs_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/min_xorsign_abs_f32.ll b/cuda-oxide/intrinsics/probes/min_xorsign_abs_f32.ll index 00d493fb9b..e5b279120d 100644 --- a/cuda-oxide/intrinsics/probes/min_xorsign_abs_f32.ll +++ b/cuda-oxide/intrinsics/probes/min_xorsign_abs_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k128_s32_b1_and_popc.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k128_s32_b1_and_popc.ll index f6b6efffc6..b4e975807f 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k128_s32_b1_and_popc.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k128_s32_b1_and_popc.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k128_s32_b1_xor_popc.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k128_s32_b1_xor_popc.ll index 2439dbb0a2..48298350b4 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k128_s32_b1_xor_popc.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k128_s32_b1_xor_popc.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_f16_f16.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_f16_f16.ll index 652608c351..1f4ac087b0 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_f16_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_f16_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_f32_bf16.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_f32_bf16.ll index 218c6a153d..5af4a8f4a4 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_f32_bf16.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_f32_bf16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_f32_f16.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_f32_f16.ll index 98a8796d1a..2c12241d56 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_f32_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_f32_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e4m3_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e4m3_e4m3.ll index ea2b871879..7f6dcd01fb 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e4m3_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e4m3_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e4m3_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e4m3_e5m2.ll index d402c9345b..1eb429205f 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e4m3_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e4m3_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e5m2_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e5m2_e4m3.ll index 012b647b84..358d4a758b 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e5m2_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e5m2_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e5m2_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e5m2_e5m2.ll index ca6b8b622e..1eb739ca07 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e5m2_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f16_e5m2_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e4m3_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e4m3_e4m3.ll index 9ac4b99413..3e8c04d5aa 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e4m3_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e4m3_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e4m3_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e4m3_e5m2.ll index fedb4958f9..6fe755274b 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e4m3_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e4m3_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e5m2_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e5m2_e4m3.ll index 49c334273a..b24fe8848f 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e5m2_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e5m2_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e5m2_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e5m2_e5m2.ll index ff4e222fba..f3457fe575 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e5m2_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_fp8_f32_e5m2_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8.ll index deeceaf268..c9e4338e39 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8_satfinite.ll index b673871a2c..57736014c8 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8_u8.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8_u8.ll index 39f0eb5f27..2e1bf439f5 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8_u8_satfinite.ll index 90ce2a048a..6eab8fcfff 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_s8_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8.ll index 95289150dd..ef431e8b6c 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8_s8.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8_s8.ll index ce48ca3fb2..dabaebed4e 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8_s8_satfinite.ll index 41575f7eb1..3086c668ab 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8_satfinite.ll index f68e9957d7..c6c3278ac8 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k16_s32_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k256_s32_b1_and_popc.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k256_s32_b1_and_popc.ll index 88fd2d1db4..6dea72bee6 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k256_s32_b1_and_popc.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k256_s32_b1_and_popc.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k256_s32_b1_xor_popc.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k256_s32_b1_xor_popc.ll index eabe1699f9..ad076de539 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k256_s32_b1_xor_popc.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k256_s32_b1_xor_popc.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e2m1.ll index bfe8f85551..5939f48bd6 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e2m3.ll index 6adfd6b3bf..d515a081cd 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e3m2.ll index e717827acd..797b7bc165 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e4m3.ll index aa9060fb70..c1095eeeda 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e5m2.ll index bf6340d5ff..f9b481b1cd 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m1_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e2m1.ll index 2509929724..f6fa8a8b94 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e2m3.ll index c3b4883df2..e591428a40 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e3m2.ll index be30a20548..336a2fc9eb 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e4m3.ll index fe3e9b05ef..a64dc83aa0 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e5m2.ll index 4ad113bf10..3de079ae34 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e2m3_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e2m1.ll index 41971701a2..1d594861a2 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e2m3.ll index e7fa742e27..bf6fb30fb5 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e3m2.ll index 997dab02e5..5157475bcb 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e4m3.ll index 318d6b23dc..926981ee3b 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e5m2.ll index 0fe8591555..bdb1cf6f7b 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e3m2_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e2m1.ll index 30f38f7304..9e4d457175 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e2m3.ll index 2c9306a880..6471db0fc1 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e3m2.ll index 06fdc0d9a6..fa2dc0f0ce 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e4m3.ll index 284c2a0c4f..145d623ae6 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e5m2.ll index 9816b24997..f66ebb38b4 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e4m3_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e2m1.ll index 3d3763d428..2e390a8530 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e2m3.ll index 3ad587c294..9297a08dc5 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e3m2.ll index e9eca1d72d..69251ec585 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e4m3.ll index d7a01bf945..2781c2f4b4 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e5m2.ll index 320ffe0f20..b35ae95c6d 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f16_e5m2_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e2m1.ll index 97d665fda6..023e957873 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e2m3.ll index ed3ea83db1..32129ec34b 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e3m2.ll index f925f54b09..16a7482476 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e4m3.ll index b84a5e1392..cf67a6dcba 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e5m2.ll index 705b43fddb..0cf24c03bc 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m1_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e2m1.ll index d02d030314..564006eda6 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e2m3.ll index cf0dc43b0b..8c02a0e129 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e3m2.ll index 5e983e32a9..d760337d89 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e4m3.ll index 73140c8270..dbf05231a8 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e5m2.ll index ee816d78a1..ad457ec485 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e2m3_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e2m1.ll index b51269d646..caf2f81c46 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e2m3.ll index 386bee2ad5..09aea1a582 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e3m2.ll index b8690c3ea1..6c8cb85b77 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e4m3.ll index 4442ad8f30..bf5d5a5c52 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e5m2.ll index 62e45327cf..579f745a35 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e3m2_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e2m1.ll index 1e02ad901b..4f44f4456d 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e2m3.ll index e00ca6a208..adf33e2d96 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e3m2.ll index bb51c964cb..4aa739cd02 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e4m3.ll index 7bc527e49c..dc244a6404 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e5m2.ll index ea329344f1..8ad1fdeae2 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e4m3_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e2m1.ll index df91d55b21..9ba9cbfe62 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e2m3.ll index 860df61829..f9baad4970 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e3m2.ll index e7ac86e1b4..64f3d442e4 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e4m3.ll index 338584e8b7..36a9e0359b 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e5m2.ll index d625a1cfc6..9c66a24879 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_f32_e5m2_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e4m3_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e4m3_e4m3.ll index e100051397..fc89e0c7f5 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e4m3_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e4m3_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e4m3_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e4m3_e5m2.ll index 006b97025d..39c7cc91f9 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e4m3_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e4m3_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e5m2_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e5m2_e4m3.ll index 9c6f5c8ffe..6b4bf4e766 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e5m2_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e5m2_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e5m2_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e5m2_e5m2.ll index 3373a3b83b..8d9e16e16b 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e5m2_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f16_e5m2_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e4m3_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e4m3_e4m3.ll index db402ed264..e3187f91a5 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e4m3_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e4m3_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e4m3_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e4m3_e5m2.ll index e347fb9514..4a58100950 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e4m3_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e4m3_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e5m2_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e5m2_e4m3.ll index ff871914d6..3f1f8acff0 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e5m2_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e5m2_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e5m2_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e5m2_e5m2.ll index c27c96cd23..4872a8a802 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e5m2_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_fp8_f32_e5m2_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e2m1.ll index f0d030a083..167fe6ae17 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e2m3.ll index 85db2eaea5..e4dab78659 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e3m2.ll index 88f9219083..fdaa9e41fc 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e4m3.ll index d89df3b11e..b20766fe3d 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e5m2.ll index 875bb88823..57cdac3a4a 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m1_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e2m1.ll index 781d451c5b..dcaaa07ec6 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e2m3.ll index 3e6e8264ba..2f40a5d5ef 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e3m2.ll index 36b28d2b94..8240045dcd 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e4m3.ll index 1a8fe73562..5eefe1b2bb 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e5m2.ll index 9e591ef10d..80f036c234 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e2m3_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e2m1.ll index d8a476cb12..6c4bc824b5 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e2m3.ll index e0d5b626b0..df076afb63 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e3m2.ll index d6782b7ff7..dc16a1f5d6 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e4m3.ll index 1dfde77274..0ebb353839 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e5m2.ll index 0dc67dcf02..e8f0537784 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e3m2_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e2m1.ll index bbf220bf75..9079e653b9 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e2m3.ll index d8646301a4..07068d359f 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e3m2.ll index c785e1e389..b474a65f08 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e4m3.ll index 4d16134d24..3fdbdcfe91 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e5m2.ll index ef0fea21e8..266e1763f0 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e4m3_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e2m1.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e2m1.ll index 587ee7805c..c7b7bd41e1 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e2m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e2m3.ll index a758bfc887..51c88be059 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e3m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e3m2.ll index 4141ae8f95..af35e5c16e 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e4m3.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e4m3.ll index be63082608..30093809ec 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e5m2.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e5m2.ll index 5200a872a0..92a028b00b 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_mxf8f6f4_f32_e5m2_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4.ll index e1acbbfc07..dac14b90e9 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4_satfinite.ll index c723dc2385..a6fe183686 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4_u4.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4_u4.ll index 62bbdb9028..2a0fc93965 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4_u4_satfinite.ll index 9ca744b0ae..926a8bb99c 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s4_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8.ll index 0504b55d29..20dc0f9d05 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8_satfinite.ll index 9416c510e7..8b065f6467 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8_u8.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8_u8.ll index c0b6d5a57c..2951a7dc3f 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8_u8_satfinite.ll index 02d6b57555..513f5799a8 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_s8_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4.ll index 3b28e3c1e1..74e4563764 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4_s4.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4_s4.ll index 008a3850b9..5f1d125b4e 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4_s4_satfinite.ll index eb24aa26ee..4f871bddf4 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4_satfinite.ll index 6f66bd0709..94021d7cb5 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8.ll index c3572d0c28..25e91fba2e 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8_s8.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8_s8.ll index ff618a524c..022d479c65 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8_s8_satfinite.ll index ebc0871a6b..bcc7415949 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8_satfinite.ll index 60ac259359..a129973805 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k32_s32_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k4_f32_tf32.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k4_f32_tf32.ll index 0f6fc9940a..e14fac55eb 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k4_f32_tf32.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k4_f32_tf32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k64_mxf4_f32.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k64_mxf4_f32.ll index 57456710a6..6f57a7a88e 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k64_mxf4_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k64_mxf4_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4.ll index b4f88f92d1..ace6a8b13b 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4_satfinite.ll index 0c58f05cb5..588f16337d 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4_u4.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4_u4.ll index a0f178a16f..789b6fb3e8 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4_u4_satfinite.ll index 26d5fdeaad..f0c7ac6419 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_s4_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4.ll index 0258fd29c9..bf3c25e223 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4_s4.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4_s4.ll index 17d2d2df9d..31152082a0 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4_s4_satfinite.ll index 68215fb8d7..6a4fb81360 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4_satfinite.ll index 1b14083042..4518dedb01 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k64_s32_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k8_f16_f16.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k8_f16_f16.ll index c2b9b0439b..81b4d5109b 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k8_f16_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k8_f16_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k8_f32_bf16.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k8_f32_bf16.ll index ebca706f54..b3610ed804 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k8_f32_bf16.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k8_f32_bf16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k8_f32_f16.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k8_f32_f16.ll index a1f6fbbaa8..81f1d5e627 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k8_f32_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k8_f32_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m16n8k8_f32_tf32.ll b/cuda-oxide/intrinsics/probes/mma_m16n8k8_f32_tf32.ll index 4b4268be6c..70d750d4f3 100644 --- a/cuda-oxide/intrinsics/probes/mma_m16n8k8_f32_tf32.ll +++ b/cuda-oxide/intrinsics/probes/mma_m16n8k8_f32_tf32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k128_s32_b1_and_popc.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k128_s32_b1_and_popc.ll index 714e46875f..161dc465f3 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k128_s32_b1_and_popc.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k128_s32_b1_and_popc.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k128_s32_b1_xor_popc.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k128_s32_b1_xor_popc.ll index 0db5ef34e7..8ac472b467 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k128_s32_b1_xor_popc.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k128_s32_b1_xor_popc.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8.ll index f5fc4dc137..2fe8f43988 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8_satfinite.ll index b7ed6ff12b..26e93b6152 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8_u8.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8_u8.ll index 07e2e40611..ad65b21bb2 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8_u8_satfinite.ll index 19973729a1..8a0de5a5b0 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_s8_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8.ll index eede881133..8554365ac2 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8_s8.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8_s8.ll index 82f449c955..b0c8f73e0d 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8_s8_satfinite.ll index e76fa19f90..cf1462184b 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8_satfinite.ll index 67234e1e24..c604c51682 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k16_s32_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4.ll index 6de3398421..b4483a2cb6 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4_satfinite.ll index 823cbf08e2..6e504a815b 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4_u4.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4_u4.ll index 814f9f7084..d239610e61 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4_u4_satfinite.ll index d2e6ca064a..5d8da01fea 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_s4_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4.ll index 312d77dcbc..0856668aa5 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4_s4.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4_s4.ll index e9b0cfb591..8122452cf0 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4_s4_satfinite.ll index 6771d74769..667c601db7 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4_satfinite.ll index 2df24153ea..01c77fe314 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k32_s32_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_m8n8k4_f64.ll b/cuda-oxide/intrinsics/probes/mma_m8n8k4_f64.ll index b535b6d1b0..c0e48af725 100644 --- a/cuda-oxide/intrinsics/probes/mma_m8n8k4_f64.ll +++ b/cuda-oxide/intrinsics/probes/mma_m8n8k4_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4.ll index 7fff20829c..4a5db338bc 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4_satfinite.ll index 81c1ff773d..165dd7bc32 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4_u4.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4_u4.ll index eac4976f7e..e241388ca7 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4_u4_satfinite.ll index dff8b73012..e24394c5cd 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_s4_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4.ll index a0e63b00a5..deef84c7e0 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4_s4.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4_s4.ll index 42834363f3..fc4d6cf6c3 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4_s4_satfinite.ll index afd7117744..83f4a2966d 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4_satfinite.ll index ad6562648e..fc597df0c0 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k128_s32_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8.ll index 2660c8ca8c..1c9257e0e6 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8_satfinite.ll index 0f92366289..ecc0698a37 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8_u8.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8_u8.ll index 33a1273e7d..3085020221 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8_u8_satfinite.ll index dd799f4a1a..8d6da3a922 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_s8_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8.ll index 5f2818f1fc..ac25f898d9 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8_s8.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8_s8.ll index b7e34a54c3..b24c7cd1d2 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8_s8_satfinite.ll index 866d336b30..0fbee4abfe 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8_satfinite.ll index 6b4e09a4a6..9736d73f1f 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k32_s32_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e4m3_e4m3_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e4m3_e4m3_f32.ll index 8f08a10bcf..99f7c07fd1 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e4m3_e4m3_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e4m3_e4m3_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e4m3_e5m2_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e4m3_e5m2_f32.ll index 93f84a0eb0..355f2b340d 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e4m3_e5m2_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e4m3_e5m2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e5m2_e4m3_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e5m2_e4m3_f32.ll index 83b41fddd4..d3071735a9 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e5m2_e4m3_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e5m2_e4m3_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e5m2_e5m2_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e5m2_e5m2_f32.ll index 2db5fbeb95..cbe45dba48 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e5m2_e5m2_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_f32_e5m2_e5m2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4.ll index 3240af04ba..38824b0e60 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4_satfinite.ll index 3f228335f8..3bb11cc4da 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4_u4.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4_u4.ll index 4a65933d7c..f13717d571 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4_u4_satfinite.ll index db6d863850..848723e429 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s4_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8.ll index 08426cefca..a8eb1b145b 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8_satfinite.ll index 1c2dcdef4c..aa5efe706f 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8_u8.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8_u8.ll index 56e8861c54..e90bd98cda 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8_u8_satfinite.ll index a8f3b233b4..90b9dd2c07 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_s8_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4.ll index 2434ceb95f..8c16c8adae 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4_s4.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4_s4.ll index 0202c7b043..51de7e7216 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4_s4_satfinite.ll index d35edf0c85..e94009ced6 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4_satfinite.ll index 6caae8d17e..0b74a031de 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8.ll index 28d0502e5b..75b2cffbba 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8_s8.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8_s8.ll index 4a58cc912d..b9aa402619 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8_s8_satfinite.ll index 30ab1dec8a..b13d189c3a 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8_satfinite.ll index b791d8021a..468bac49ff 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_m16n8k64_s32_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4.ll index 4c66c98ea1..da3081ca09 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4_satfinite.ll index 6bacc4b46d..a69df09f51 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4_u4.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4_u4.ll index 28a6aff9c7..269d1f5815 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4_u4_satfinite.ll index bddabe63f6..8fd43f2c1b 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_s4_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4.ll index 4275d5994f..cd7e6fde06 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4_s4.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4_s4.ll index 89072de7dc..e4a9e5b57c 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4_s4_satfinite.ll index 68d78a9623..c15236af28 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4_satfinite.ll index c2f987a6ef..4bfa50ceb4 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k128_s32_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f16_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f16_f16.ll index bca11fdff5..d2ababde42 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f16_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f16_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f32_bf16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f32_bf16.ll index d916e14b50..001d8e5e7d 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f32_bf16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f32_bf16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f32_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f32_f16.ll index 72bb48cc64..347f83132f 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f32_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f32_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f32_tf32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f32_tf32.ll index 0c265b70dd..d7359fa5b6 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f32_tf32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k16_f32_tf32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_f16_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_f16_f16.ll index ce98defdde..5ea3085913 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_f16_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_f16_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_f32_bf16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_f32_bf16.ll index faffdd3967..aa0624cebe 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_f32_bf16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_f32_bf16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_f32_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_f32_f16.ll index 4bd31db113..35de479a75 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_f32_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_f32_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8.ll index 4b074efffc..d34b135f89 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8_satfinite.ll index 33f49f21c2..eee3b1d195 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8_u8.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8_u8.ll index cab86d8bbe..f5d367b679 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8_u8_satfinite.ll index 7eda38083e..83cb827c40 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_s8_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8.ll index 7587fc6021..6b5dddeb3a 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8_s8.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8_s8.ll index ba13859c26..bec63db85b 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8_s8_satfinite.ll index 22b712fb30..da80a8c397 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8_satfinite.ll index 4ba7604b1c..4c1cc2358c 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k32_s32_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16.ll index b2e3f27d93..fcd266c91b 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m3_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m3_f16.ll index c9dee2d132..48bc7d62a0 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m3_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m3_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e3m2_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e3m2_f16.ll index e4ab24f478..8392cbd552 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e3m2_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e3m2_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e4m3_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e4m3_f16.ll index 13becb1072..8af461e520 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e4m3_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e4m3_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e5m2_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e5m2_f16.ll index ab8fb9ec61..fe19f3b945 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e5m2_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e5m2_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e2m1_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e2m1_f16.ll index 6551438c19..eeddf43789 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e2m1_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e2m1_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e2m3_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e2m3_f16.ll index 6d3f164cc3..b300d70447 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e2m3_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e2m3_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e3m2_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e3m2_f16.ll index f865cb6f43..7ab9664953 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e3m2_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e3m2_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e4m3_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e4m3_f16.ll index 219502ee1f..a95a58290a 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e4m3_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e4m3_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e5m2_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e5m2_f16.ll index 0c49cabee5..f450d9b962 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e5m2_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m3_e5m2_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e2m1_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e2m1_f16.ll index 1d4e67f597..626d50de66 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e2m1_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e2m1_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e2m3_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e2m3_f16.ll index d8c4b60004..77aff0e5be 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e2m3_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e2m3_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e3m2_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e3m2_f16.ll index 7a98060465..bc97a19224 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e3m2_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e3m2_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e4m3_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e4m3_f16.ll index 82106a6efe..491c10f87f 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e4m3_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e4m3_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e5m2_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e5m2_f16.ll index 6cc01ed942..7ac5b9ab06 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e5m2_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e3m2_e5m2_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e2m1_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e2m1_f16.ll index f297a73933..38d6af5f91 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e2m1_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e2m1_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e2m3_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e2m3_f16.ll index de8d838d51..4963dc0c52 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e2m3_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e2m3_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e3m2_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e3m2_f16.ll index 8efe25ffaa..58952f022d 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e3m2_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e3m2_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e4m3_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e4m3_f16.ll index 9ef8191ef4..0bc0fd29d4 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e4m3_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e4m3_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e5m2_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e5m2_f16.ll index 4e6ed12ede..482ed78e27 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e5m2_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e4m3_e5m2_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e2m1_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e2m1_f16.ll index 1418484dd4..bffb2d93cc 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e2m1_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e2m1_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e2m3_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e2m3_f16.ll index 5263d6e447..91ff6c0161 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e2m3_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e2m3_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e3m2_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e3m2_f16.ll index 6c2656d444..8281857b9f 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e3m2_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e3m2_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e4m3_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e4m3_f16.ll index 259aacad7f..3a354179a8 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e4m3_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e4m3_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e5m2_f16.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e5m2_f16.ll index 51329ce5b5..749bb0ccef 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e5m2_f16.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e5m2_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m1_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m1_f32.ll index 95d6bd77a4..c6427f895e 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m1_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m1_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m3_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m3_f32.ll index 29895bd722..7275ba6a25 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m3_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m3_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e3m2_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e3m2_f32.ll index 2c7799f2b5..93ba1412e6 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e3m2_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e3m2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e4m3_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e4m3_f32.ll index 09313b83b6..77dcebfc8d 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e4m3_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e4m3_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e5m2_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e5m2_f32.ll index 61c86d8d87..295c1b69f8 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e5m2_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e5m2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e2m1_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e2m1_f32.ll index 6738701767..0c907e231b 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e2m1_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e2m1_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e2m3_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e2m3_f32.ll index ca17c54e27..2d0ef268ea 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e2m3_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e2m3_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e3m2_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e3m2_f32.ll index d4c9ca79ff..79f9ac3814 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e3m2_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e3m2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e4m3_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e4m3_f32.ll index 8394efd362..cdcd6c9cce 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e4m3_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e4m3_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e5m2_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e5m2_f32.ll index afc387c12b..5b7c1883f8 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e5m2_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m3_e5m2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e2m1_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e2m1_f32.ll index 22ae36e69d..2bcff37597 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e2m1_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e2m1_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e2m3_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e2m3_f32.ll index 248875a865..a353ca83c9 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e2m3_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e2m3_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e3m2_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e3m2_f32.ll index c2960c1668..1403814983 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e3m2_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e3m2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e4m3_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e4m3_f32.ll index 2ccf2d7b04..97979374f2 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e4m3_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e4m3_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e5m2_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e5m2_f32.ll index f308d53123..d3570c9d50 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e5m2_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e3m2_e5m2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e2m1_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e2m1_f32.ll index ecd01a64e9..a9745e31db 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e2m1_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e2m1_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e2m3_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e2m3_f32.ll index 269527c06c..5b7234d279 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e2m3_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e2m3_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e3m2_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e3m2_f32.ll index 8e6ae98aed..9fe78224aa 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e3m2_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e3m2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e4m3_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e4m3_f32.ll index 40bf39d8c5..41bb0a5303 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e4m3_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e4m3_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e5m2_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e5m2_f32.ll index c4763a9ba4..78b2e6d42a 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e5m2_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e4m3_e5m2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e2m1_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e2m1_f32.ll index f6359985c2..b84572ba8b 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e2m1_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e2m1_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e2m3_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e2m3_f32.ll index 0bcde6d88d..9db665bc6a 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e2m3_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e2m3_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e3m2_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e3m2_f32.ll index db0adee7da..6e5561b592 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e3m2_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e3m2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e4m3_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e4m3_f32.ll index 9863082e13..5f3b96b3e2 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e4m3_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e4m3_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e5m2_f32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e5m2_f32.ll index 3899adf96f..b440e60251 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e5m2_f32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e5m2_e5m2_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4.ll index 31a8def2ec..6debe9a853 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4_satfinite.ll index a72569eed2..3fd76b2cae 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4_u4.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4_u4.ll index 674461350f..ae86ac4772 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4_u4_satfinite.ll index d118cad5c3..6873f0d2ee 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s4_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8.ll index 79e84349e3..c84a84d085 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8_satfinite.ll index 84ccf7775c..dd105457ef 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8_u8.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8_u8.ll index 2fa6cfe29b..a20c93818c 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8_u8_satfinite.ll index 38133e215d..d4895fe1e8 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_s8_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4.ll index 4377995370..5160166e76 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4_s4.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4_s4.ll index 12a9a1e6e3..2bfe38e4df 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4_s4.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4_s4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4_s4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4_s4_satfinite.ll index cc3cafe88a..51eca5f98d 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4_s4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4_s4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4_satfinite.ll index 8a207480cc..af7b9077e9 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u4_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8.ll index 069921a564..d4a69d7c17 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8_s8.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8_s8.ll index ecf955a115..20c4d38100 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8_s8.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8_s8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8_s8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8_s8_satfinite.ll index 7401704d55..0441b0e264 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8_s8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8_s8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8_satfinite.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8_satfinite.ll index cc4982a40e..12c6fab0c1 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8_satfinite.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k64_s32_u8_satfinite.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k8_f32_tf32.ll b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k8_f32_tf32.ll index a98c36fef4..f1dda597df 100644 --- a/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k8_f32_tf32.ll +++ b/cuda-oxide/intrinsics/probes/mma_sp_ordered_metadata_m16n8k8_f32_tf32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/movmatrix_trans_b16.ll b/cuda-oxide/intrinsics/probes/movmatrix_trans_b16.ll index f18fabb281..57f19ff132 100644 --- a/cuda-oxide/intrinsics/probes/movmatrix_trans_b16.ll +++ b/cuda-oxide/intrinsics/probes/movmatrix_trans_b16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_bf16x2.ll b/cuda-oxide/intrinsics/probes/mul_bf16x2.ll index 936b3e8f54..c8f4f5cf77 100644 --- a/cuda-oxide/intrinsics/probes/mul_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/mul_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_f16x2.ll b/cuda-oxide/intrinsics/probes/mul_f16x2.ll index 4d1be55d11..e0dbd8e540 100644 --- a/cuda-oxide/intrinsics/probes/mul_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/mul_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_f32x2.ll b/cuda-oxide/intrinsics/probes/mul_f32x2.ll index e1051c262b..c0a823b673 100644 --- a/cuda-oxide/intrinsics/probes/mul_f32x2.ll +++ b/cuda-oxide/intrinsics/probes/mul_f32x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_ftz_f32x2.ll b/cuda-oxide/intrinsics/probes/mul_ftz_f32x2.ll index eb3954376c..6a6993db59 100644 --- a/cuda-oxide/intrinsics/probes/mul_ftz_f32x2.ll +++ b/cuda-oxide/intrinsics/probes/mul_ftz_f32x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_rm_f32.ll b/cuda-oxide/intrinsics/probes/mul_rm_f32.ll index 1f0f85ed45..d0fce3530c 100644 --- a/cuda-oxide/intrinsics/probes/mul_rm_f32.ll +++ b/cuda-oxide/intrinsics/probes/mul_rm_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_rm_f64.ll b/cuda-oxide/intrinsics/probes/mul_rm_f64.ll index ae16860107..34d0b3910d 100644 --- a/cuda-oxide/intrinsics/probes/mul_rm_f64.ll +++ b/cuda-oxide/intrinsics/probes/mul_rm_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_rm_ftz_f32.ll b/cuda-oxide/intrinsics/probes/mul_rm_ftz_f32.ll index 799c0671b8..f6360122e3 100644 --- a/cuda-oxide/intrinsics/probes/mul_rm_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/mul_rm_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_rn_f32.ll b/cuda-oxide/intrinsics/probes/mul_rn_f32.ll index 8a39d1826c..c79975c6d5 100644 --- a/cuda-oxide/intrinsics/probes/mul_rn_f32.ll +++ b/cuda-oxide/intrinsics/probes/mul_rn_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_rn_f64.ll b/cuda-oxide/intrinsics/probes/mul_rn_f64.ll index f2b714e40a..06e39a09b1 100644 --- a/cuda-oxide/intrinsics/probes/mul_rn_f64.ll +++ b/cuda-oxide/intrinsics/probes/mul_rn_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_rn_ftz_f32.ll b/cuda-oxide/intrinsics/probes/mul_rn_ftz_f32.ll index 4dd4ed6d9e..d8bd06210c 100644 --- a/cuda-oxide/intrinsics/probes/mul_rn_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/mul_rn_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_rp_f32.ll b/cuda-oxide/intrinsics/probes/mul_rp_f32.ll index f87a7c919f..38e6693b2d 100644 --- a/cuda-oxide/intrinsics/probes/mul_rp_f32.ll +++ b/cuda-oxide/intrinsics/probes/mul_rp_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_rp_f64.ll b/cuda-oxide/intrinsics/probes/mul_rp_f64.ll index a52de24599..c5de40e65f 100644 --- a/cuda-oxide/intrinsics/probes/mul_rp_f64.ll +++ b/cuda-oxide/intrinsics/probes/mul_rp_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_rp_ftz_f32.ll b/cuda-oxide/intrinsics/probes/mul_rp_ftz_f32.ll index d09969bbb1..9d7aa75db6 100644 --- a/cuda-oxide/intrinsics/probes/mul_rp_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/mul_rp_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_rz_f32.ll b/cuda-oxide/intrinsics/probes/mul_rz_f32.ll index f0bb681b6c..4ae401fa61 100644 --- a/cuda-oxide/intrinsics/probes/mul_rz_f32.ll +++ b/cuda-oxide/intrinsics/probes/mul_rz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_rz_f64.ll b/cuda-oxide/intrinsics/probes/mul_rz_f64.ll index 8cefde1cda..7b0078b2a2 100644 --- a/cuda-oxide/intrinsics/probes/mul_rz_f64.ll +++ b/cuda-oxide/intrinsics/probes/mul_rz_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/mul_rz_ftz_f32.ll b/cuda-oxide/intrinsics/probes/mul_rz_ftz_f32.ll index 70d9c7ed2a..3c57f38cbd 100644 --- a/cuda-oxide/intrinsics/probes/mul_rz_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/mul_rz_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/nanosleep.ll b/cuda-oxide/intrinsics/probes/nanosleep.ll index d6ed6b04f8..d517edece9 100644 --- a/cuda-oxide/intrinsics/probes/nanosleep.ll +++ b/cuda-oxide/intrinsics/probes/nanosleep.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/neg_bf16x2.ll b/cuda-oxide/intrinsics/probes/neg_bf16x2.ll index b08bee3172..cf8849ef1c 100644 --- a/cuda-oxide/intrinsics/probes/neg_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/neg_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/neg_f16x2.ll b/cuda-oxide/intrinsics/probes/neg_f16x2.ll index 329c7f8d4e..5a2c02e94d 100644 --- a/cuda-oxide/intrinsics/probes/neg_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/neg_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/nsmid.lib_nvvm.ll b/cuda-oxide/intrinsics/probes/nsmid.lib_nvvm.ll index 683964c895..98f058bfea 100644 --- a/cuda-oxide/intrinsics/probes/nsmid.lib_nvvm.ll +++ b/cuda-oxide/intrinsics/probes/nsmid.lib_nvvm.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/nsmid.llvm_nvptx.ll b/cuda-oxide/intrinsics/probes/nsmid.llvm_nvptx.ll index 0b791a955c..36306a20c5 100644 --- a/cuda-oxide/intrinsics/probes/nsmid.llvm_nvptx.ll +++ b/cuda-oxide/intrinsics/probes/nsmid.llvm_nvptx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/nwarpid.lib_nvvm.ll b/cuda-oxide/intrinsics/probes/nwarpid.lib_nvvm.ll index 9a3b9c71de..8210d37cdf 100644 --- a/cuda-oxide/intrinsics/probes/nwarpid.lib_nvvm.ll +++ b/cuda-oxide/intrinsics/probes/nwarpid.lib_nvvm.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/nwarpid.llvm_nvptx.ll b/cuda-oxide/intrinsics/probes/nwarpid.llvm_nvptx.ll index d76519dc21..aa045c4c99 100644 --- a/cuda-oxide/intrinsics/probes/nwarpid.llvm_nvptx.ll +++ b/cuda-oxide/intrinsics/probes/nwarpid.llvm_nvptx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/packed_atomic_add_bf16x2.ll b/cuda-oxide/intrinsics/probes/packed_atomic_add_bf16x2.ll index 7fb753bb40..9556f56e7c 100644 --- a/cuda-oxide/intrinsics/probes/packed_atomic_add_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/packed_atomic_add_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/packed_atomic_add_f16x2.ll b/cuda-oxide/intrinsics/probes/packed_atomic_add_f16x2.ll index 2af43904e0..86b3d9c598 100644 --- a/cuda-oxide/intrinsics/probes/packed_atomic_add_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/packed_atomic_add_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/pmevent.ll b/cuda-oxide/intrinsics/probes/pmevent.ll index a7a611d11f..ab3dab3c77 100644 --- a/cuda-oxide/intrinsics/probes/pmevent.ll +++ b/cuda-oxide/intrinsics/probes/pmevent.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/prefetch_tma_descriptor.ll b/cuda-oxide/intrinsics/probes/prefetch_tma_descriptor.ll index aaf83cc626..044d021c04 100644 --- a/cuda-oxide/intrinsics/probes/prefetch_tma_descriptor.ll +++ b/cuda-oxide/intrinsics/probes/prefetch_tma_descriptor.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/prmt.ll b/cuda-oxide/intrinsics/probes/prmt.ll index 4192045777..c7fc1df939 100644 --- a/cuda-oxide/intrinsics/probes/prmt.ll +++ b/cuda-oxide/intrinsics/probes/prmt.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/prmt_b4e.ll b/cuda-oxide/intrinsics/probes/prmt_b4e.ll index 9edf78eb9a..036a16b449 100644 --- a/cuda-oxide/intrinsics/probes/prmt_b4e.ll +++ b/cuda-oxide/intrinsics/probes/prmt_b4e.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/prmt_ecl.ll b/cuda-oxide/intrinsics/probes/prmt_ecl.ll index 708bc37d98..f6661f0499 100644 --- a/cuda-oxide/intrinsics/probes/prmt_ecl.ll +++ b/cuda-oxide/intrinsics/probes/prmt_ecl.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/prmt_ecr.ll b/cuda-oxide/intrinsics/probes/prmt_ecr.ll index ad6d026187..97169d0e0c 100644 --- a/cuda-oxide/intrinsics/probes/prmt_ecr.ll +++ b/cuda-oxide/intrinsics/probes/prmt_ecr.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/prmt_f4e.ll b/cuda-oxide/intrinsics/probes/prmt_f4e.ll index 5a486f2893..50a0f630c2 100644 --- a/cuda-oxide/intrinsics/probes/prmt_f4e.ll +++ b/cuda-oxide/intrinsics/probes/prmt_f4e.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/prmt_rc16.ll b/cuda-oxide/intrinsics/probes/prmt_rc16.ll index 5758905650..7ce56e4461 100644 --- a/cuda-oxide/intrinsics/probes/prmt_rc16.ll +++ b/cuda-oxide/intrinsics/probes/prmt_rc16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/prmt_rc8.ll b/cuda-oxide/intrinsics/probes/prmt_rc8.ll index 285290fa91..918177de94 100644 --- a/cuda-oxide/intrinsics/probes/prmt_rc8.ll +++ b/cuda-oxide/intrinsics/probes/prmt_rc8.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_approx_ftz_f32.ll b/cuda-oxide/intrinsics/probes/rcp_approx_ftz_f32.ll index 3668e565e8..d35fa718b4 100644 --- a/cuda-oxide/intrinsics/probes/rcp_approx_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/rcp_approx_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_approx_ftz_f64.ll b/cuda-oxide/intrinsics/probes/rcp_approx_ftz_f64.ll index cec5a0dae5..f5b08dd3e8 100644 --- a/cuda-oxide/intrinsics/probes/rcp_approx_ftz_f64.ll +++ b/cuda-oxide/intrinsics/probes/rcp_approx_ftz_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_rm_f32.ll b/cuda-oxide/intrinsics/probes/rcp_rm_f32.ll index da54467a6f..c274353bcf 100644 --- a/cuda-oxide/intrinsics/probes/rcp_rm_f32.ll +++ b/cuda-oxide/intrinsics/probes/rcp_rm_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_rm_f64.ll b/cuda-oxide/intrinsics/probes/rcp_rm_f64.ll index a4c5d9740a..a0c74dce0b 100644 --- a/cuda-oxide/intrinsics/probes/rcp_rm_f64.ll +++ b/cuda-oxide/intrinsics/probes/rcp_rm_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_rm_ftz_f32.ll b/cuda-oxide/intrinsics/probes/rcp_rm_ftz_f32.ll index 88d349a241..0250ea771b 100644 --- a/cuda-oxide/intrinsics/probes/rcp_rm_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/rcp_rm_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_rn_f32.ll b/cuda-oxide/intrinsics/probes/rcp_rn_f32.ll index bc84cf309e..36474c6f56 100644 --- a/cuda-oxide/intrinsics/probes/rcp_rn_f32.ll +++ b/cuda-oxide/intrinsics/probes/rcp_rn_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_rn_f64.ll b/cuda-oxide/intrinsics/probes/rcp_rn_f64.ll index c438c43725..c28b5b6b82 100644 --- a/cuda-oxide/intrinsics/probes/rcp_rn_f64.ll +++ b/cuda-oxide/intrinsics/probes/rcp_rn_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_rn_ftz_f32.ll b/cuda-oxide/intrinsics/probes/rcp_rn_ftz_f32.ll index 82a08f1e9f..1e45d9ca88 100644 --- a/cuda-oxide/intrinsics/probes/rcp_rn_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/rcp_rn_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_rp_f32.ll b/cuda-oxide/intrinsics/probes/rcp_rp_f32.ll index 55a731aec4..14855d149a 100644 --- a/cuda-oxide/intrinsics/probes/rcp_rp_f32.ll +++ b/cuda-oxide/intrinsics/probes/rcp_rp_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_rp_f64.ll b/cuda-oxide/intrinsics/probes/rcp_rp_f64.ll index 534b8a5b42..9478f1ca04 100644 --- a/cuda-oxide/intrinsics/probes/rcp_rp_f64.ll +++ b/cuda-oxide/intrinsics/probes/rcp_rp_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_rp_ftz_f32.ll b/cuda-oxide/intrinsics/probes/rcp_rp_ftz_f32.ll index e8731ceb27..9ce4cee816 100644 --- a/cuda-oxide/intrinsics/probes/rcp_rp_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/rcp_rp_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_rz_f32.ll b/cuda-oxide/intrinsics/probes/rcp_rz_f32.ll index 53398db66a..639bf80353 100644 --- a/cuda-oxide/intrinsics/probes/rcp_rz_f32.ll +++ b/cuda-oxide/intrinsics/probes/rcp_rz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_rz_f64.ll b/cuda-oxide/intrinsics/probes/rcp_rz_f64.ll index 5bf264b1b2..e2256d8196 100644 --- a/cuda-oxide/intrinsics/probes/rcp_rz_f64.ll +++ b/cuda-oxide/intrinsics/probes/rcp_rz_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rcp_rz_ftz_f32.ll b/cuda-oxide/intrinsics/probes/rcp_rz_ftz_f32.ll index 4d60e4c8df..86af7310c6 100644 --- a/cuda-oxide/intrinsics/probes/rcp_rz_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/rcp_rz_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_add.ll b/cuda-oxide/intrinsics/probes/redux_sync_add.ll index dfbef212d0..837aac029e 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_add.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_add.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_and.ll b/cuda-oxide/intrinsics/probes/redux_sync_and.ll index 96aaea7d80..48c761eb45 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_and.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_and.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_max_abs_f32.ll b/cuda-oxide/intrinsics/probes/redux_sync_max_abs_f32.ll index f0447b6fc4..bd04176a23 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_max_abs_f32.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_max_abs_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_max_abs_nan_f32.ll b/cuda-oxide/intrinsics/probes/redux_sync_max_abs_nan_f32.ll index c38cf66790..678a048dc8 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_max_abs_nan_f32.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_max_abs_nan_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_max_f32.ll b/cuda-oxide/intrinsics/probes/redux_sync_max_f32.ll index 027d39f14b..8a56fc9860 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_max_f32.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_max_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_max_i32.ll b/cuda-oxide/intrinsics/probes/redux_sync_max_i32.ll index 3b70967d24..121a125e5b 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_max_i32.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_max_i32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_max_nan_f32.ll b/cuda-oxide/intrinsics/probes/redux_sync_max_nan_f32.ll index 8adbc17816..55dea3ec43 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_max_nan_f32.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_max_nan_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_max_u32.ll b/cuda-oxide/intrinsics/probes/redux_sync_max_u32.ll index 66e5699e94..011951a5fc 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_max_u32.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_max_u32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_min_abs_f32.ll b/cuda-oxide/intrinsics/probes/redux_sync_min_abs_f32.ll index 24dd4b363d..5ee5a0096b 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_min_abs_f32.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_min_abs_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_min_abs_nan_f32.ll b/cuda-oxide/intrinsics/probes/redux_sync_min_abs_nan_f32.ll index f2f1fee68a..dd27be68cf 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_min_abs_nan_f32.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_min_abs_nan_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_min_f32.ll b/cuda-oxide/intrinsics/probes/redux_sync_min_f32.ll index 4a9513b413..d550cdb3f5 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_min_f32.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_min_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_min_i32.ll b/cuda-oxide/intrinsics/probes/redux_sync_min_i32.ll index de361f3d3d..eeb9afb305 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_min_i32.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_min_i32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_min_nan_f32.ll b/cuda-oxide/intrinsics/probes/redux_sync_min_nan_f32.ll index df8a221cec..6d69a12a5f 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_min_nan_f32.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_min_nan_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_min_u32.ll b/cuda-oxide/intrinsics/probes/redux_sync_min_u32.ll index ba251716f9..f7e49b550f 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_min_u32.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_min_u32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_or.ll b/cuda-oxide/intrinsics/probes/redux_sync_or.ll index 71c3d4949a..0e5cfe73ff 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_or.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_or.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/redux_sync_xor.ll b/cuda-oxide/intrinsics/probes/redux_sync_xor.ll index b86901de48..5239bb0c43 100644 --- a/cuda-oxide/intrinsics/probes/redux_sync_xor.ll +++ b/cuda-oxide/intrinsics/probes/redux_sync_xor.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rsqrt_approx_f32.ll b/cuda-oxide/intrinsics/probes/rsqrt_approx_f32.ll index 79c4ac55a1..1f7bf48483 100644 --- a/cuda-oxide/intrinsics/probes/rsqrt_approx_f32.ll +++ b/cuda-oxide/intrinsics/probes/rsqrt_approx_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rsqrt_approx_f64.ll b/cuda-oxide/intrinsics/probes/rsqrt_approx_f64.ll index 2030fcbf8a..3fec74a70f 100644 --- a/cuda-oxide/intrinsics/probes/rsqrt_approx_f64.ll +++ b/cuda-oxide/intrinsics/probes/rsqrt_approx_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/rsqrt_approx_ftz_f32.ll b/cuda-oxide/intrinsics/probes/rsqrt_approx_ftz_f32.ll index 9f59023b62..db8696145b 100644 --- a/cuda-oxide/intrinsics/probes/rsqrt_approx_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/rsqrt_approx_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/setmaxnreg_dec.ll b/cuda-oxide/intrinsics/probes/setmaxnreg_dec.ll index 6b0f5f2727..5c22fe65b9 100644 --- a/cuda-oxide/intrinsics/probes/setmaxnreg_dec.ll +++ b/cuda-oxide/intrinsics/probes/setmaxnreg_dec.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/setmaxnreg_inc.ll b/cuda-oxide/intrinsics/probes/setmaxnreg_inc.ll index 6c76580c17..95e6da88a2 100644 --- a/cuda-oxide/intrinsics/probes/setmaxnreg_inc.ll +++ b/cuda-oxide/intrinsics/probes/setmaxnreg_inc.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/shuffle_down_f32_sync.ll b/cuda-oxide/intrinsics/probes/shuffle_down_f32_sync.ll index f23400cb93..cbffa3208a 100644 --- a/cuda-oxide/intrinsics/probes/shuffle_down_f32_sync.ll +++ b/cuda-oxide/intrinsics/probes/shuffle_down_f32_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/shuffle_down_sync.ll b/cuda-oxide/intrinsics/probes/shuffle_down_sync.ll index b804447fef..b8dfade604 100644 --- a/cuda-oxide/intrinsics/probes/shuffle_down_sync.ll +++ b/cuda-oxide/intrinsics/probes/shuffle_down_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/shuffle_down_u64_sync.ll b/cuda-oxide/intrinsics/probes/shuffle_down_u64_sync.ll index 7b2a53366c..cde67975a4 100644 --- a/cuda-oxide/intrinsics/probes/shuffle_down_u64_sync.ll +++ b/cuda-oxide/intrinsics/probes/shuffle_down_u64_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/shuffle_f32_sync.ll b/cuda-oxide/intrinsics/probes/shuffle_f32_sync.ll index 9ed163eb40..393fe3d87b 100644 --- a/cuda-oxide/intrinsics/probes/shuffle_f32_sync.ll +++ b/cuda-oxide/intrinsics/probes/shuffle_f32_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/shuffle_sync.ll b/cuda-oxide/intrinsics/probes/shuffle_sync.ll index 3149630d84..f18683ca51 100644 --- a/cuda-oxide/intrinsics/probes/shuffle_sync.ll +++ b/cuda-oxide/intrinsics/probes/shuffle_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/shuffle_u64_sync.ll b/cuda-oxide/intrinsics/probes/shuffle_u64_sync.ll index 9e326cea74..c0b8635858 100644 --- a/cuda-oxide/intrinsics/probes/shuffle_u64_sync.ll +++ b/cuda-oxide/intrinsics/probes/shuffle_u64_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/shuffle_up_f32_sync.ll b/cuda-oxide/intrinsics/probes/shuffle_up_f32_sync.ll index 08ee509777..b3fe240889 100644 --- a/cuda-oxide/intrinsics/probes/shuffle_up_f32_sync.ll +++ b/cuda-oxide/intrinsics/probes/shuffle_up_f32_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/shuffle_up_sync.ll b/cuda-oxide/intrinsics/probes/shuffle_up_sync.ll index ff4afe701b..12a711e060 100644 --- a/cuda-oxide/intrinsics/probes/shuffle_up_sync.ll +++ b/cuda-oxide/intrinsics/probes/shuffle_up_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/shuffle_up_u64_sync.ll b/cuda-oxide/intrinsics/probes/shuffle_up_u64_sync.ll index bf2416179e..f7ed95935a 100644 --- a/cuda-oxide/intrinsics/probes/shuffle_up_u64_sync.ll +++ b/cuda-oxide/intrinsics/probes/shuffle_up_u64_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/shuffle_xor_f32_sync.ll b/cuda-oxide/intrinsics/probes/shuffle_xor_f32_sync.ll index 627df49b6b..fb5b4d1b80 100644 --- a/cuda-oxide/intrinsics/probes/shuffle_xor_f32_sync.ll +++ b/cuda-oxide/intrinsics/probes/shuffle_xor_f32_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/shuffle_xor_sync.ll b/cuda-oxide/intrinsics/probes/shuffle_xor_sync.ll index 090deb711f..e66345d5bb 100644 --- a/cuda-oxide/intrinsics/probes/shuffle_xor_sync.ll +++ b/cuda-oxide/intrinsics/probes/shuffle_xor_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/shuffle_xor_u64_sync.ll b/cuda-oxide/intrinsics/probes/shuffle_xor_u64_sync.ll index d1f55cb353..572c2524c2 100644 --- a/cuda-oxide/intrinsics/probes/shuffle_xor_u64_sync.ll +++ b/cuda-oxide/intrinsics/probes/shuffle_xor_u64_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sin_approx_f32.ll b/cuda-oxide/intrinsics/probes/sin_approx_f32.ll index ccb7adb83b..d91b8f99cd 100644 --- a/cuda-oxide/intrinsics/probes/sin_approx_f32.ll +++ b/cuda-oxide/intrinsics/probes/sin_approx_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sin_approx_ftz_f32.ll b/cuda-oxide/intrinsics/probes/sin_approx_ftz_f32.ll index b1ae865e5e..198360cdfc 100644 --- a/cuda-oxide/intrinsics/probes/sin_approx_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/sin_approx_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/smid.lib_nvvm.ll b/cuda-oxide/intrinsics/probes/smid.lib_nvvm.ll index 27cc6b741b..870f6b79f4 100644 --- a/cuda-oxide/intrinsics/probes/smid.lib_nvvm.ll +++ b/cuda-oxide/intrinsics/probes/smid.lib_nvvm.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/smid.llvm_nvptx.ll b/cuda-oxide/intrinsics/probes/smid.llvm_nvptx.ll index e1d1f23e58..ab20d38942 100644 --- a/cuda-oxide/intrinsics/probes/smid.llvm_nvptx.ll +++ b/cuda-oxide/intrinsics/probes/smid.llvm_nvptx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_approx_f32.ll b/cuda-oxide/intrinsics/probes/sqrt_approx_f32.ll index a6c690fa4a..26510936d4 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_approx_f32.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_approx_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_approx_ftz_f32.ll b/cuda-oxide/intrinsics/probes/sqrt_approx_ftz_f32.ll index e8c5690ce3..22ddfcfc20 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_approx_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_approx_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_rm_f32.ll b/cuda-oxide/intrinsics/probes/sqrt_rm_f32.ll index 2fd7aaba7b..bb8fc99f40 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_rm_f32.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_rm_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_rm_f64.ll b/cuda-oxide/intrinsics/probes/sqrt_rm_f64.ll index 490adc4b07..cdbba0caed 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_rm_f64.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_rm_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_rm_ftz_f32.ll b/cuda-oxide/intrinsics/probes/sqrt_rm_ftz_f32.ll index 032447314a..1872a1307b 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_rm_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_rm_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_rn_f32.ll b/cuda-oxide/intrinsics/probes/sqrt_rn_f32.ll index 1cb9483460..1b1ceaeca9 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_rn_f32.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_rn_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_rn_f64.ll b/cuda-oxide/intrinsics/probes/sqrt_rn_f64.ll index b2db2a3af4..7cc0449fdf 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_rn_f64.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_rn_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_rn_ftz_f32.ll b/cuda-oxide/intrinsics/probes/sqrt_rn_ftz_f32.ll index f965981d58..86b71fc89d 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_rn_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_rn_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_rp_f32.ll b/cuda-oxide/intrinsics/probes/sqrt_rp_f32.ll index f006e4b778..be94d76fe6 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_rp_f32.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_rp_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_rp_f64.ll b/cuda-oxide/intrinsics/probes/sqrt_rp_f64.ll index 2eb9ee656a..856e26d76c 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_rp_f64.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_rp_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_rp_ftz_f32.ll b/cuda-oxide/intrinsics/probes/sqrt_rp_ftz_f32.ll index c345537036..d7200f381b 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_rp_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_rp_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_rz_f32.ll b/cuda-oxide/intrinsics/probes/sqrt_rz_f32.ll index 47c543ff0e..4c19e0002e 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_rz_f32.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_rz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_rz_f64.ll b/cuda-oxide/intrinsics/probes/sqrt_rz_f64.ll index 8a619642c1..163ce3e37e 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_rz_f64.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_rz_f64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sqrt_rz_ftz_f32.ll b/cuda-oxide/intrinsics/probes/sqrt_rz_ftz_f32.ll index d2ba34a117..c075e5ad19 100644 --- a/cuda-oxide/intrinsics/probes/sqrt_rz_ftz_f32.ll +++ b/cuda-oxide/intrinsics/probes/sqrt_rz_ftz_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x2_b16.ll b/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x2_b16.ll index 37cfa64bef..5d65185524 100644 --- a/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x2_b16.ll +++ b/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x2_b16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x2_trans_b16.ll b/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x2_trans_b16.ll index 6d2f10ebab..7a0e2c1a77 100644 --- a/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x2_trans_b16.ll +++ b/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x2_trans_b16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x4_b16.ll b/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x4_b16.ll index 18b40eb932..8d2069afbf 100644 --- a/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x4_b16.ll +++ b/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x4_b16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x4_trans_b16.ll b/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x4_trans_b16.ll index 5f28fcc9a4..cbe9570316 100644 --- a/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x4_trans_b16.ll +++ b/cuda-oxide/intrinsics/probes/stmatrix_m8n8_x4_trans_b16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sub_bf16x2.ll b/cuda-oxide/intrinsics/probes/sub_bf16x2.ll index a4bc202b9c..2748353a05 100644 --- a/cuda-oxide/intrinsics/probes/sub_bf16x2.ll +++ b/cuda-oxide/intrinsics/probes/sub_bf16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sub_f16x2.ll b/cuda-oxide/intrinsics/probes/sub_f16x2.ll index b6c92252ef..9e5cb2ed2c 100644 --- a/cuda-oxide/intrinsics/probes/sub_f16x2.ll +++ b/cuda-oxide/intrinsics/probes/sub_f16x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sub_f32x2.ll b/cuda-oxide/intrinsics/probes/sub_f32x2.ll index 2c771fe2fa..6a5ec05754 100644 --- a/cuda-oxide/intrinsics/probes/sub_f32x2.ll +++ b/cuda-oxide/intrinsics/probes/sub_f32x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sub_ftz_f32x2.ll b/cuda-oxide/intrinsics/probes/sub_ftz_f32x2.ll index 4169fe7086..1cd5597481 100644 --- a/cuda-oxide/intrinsics/probes/sub_ftz_f32x2.ll +++ b/cuda-oxide/intrinsics/probes/sub_ftz_f32x2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sync_mask.ll b/cuda-oxide/intrinsics/probes/sync_mask.ll index 165652b0ab..f69bcf22ea 100644 --- a/cuda-oxide/intrinsics/probes/sync_mask.ll +++ b/cuda-oxide/intrinsics/probes/sync_mask.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/sync_threads.ll b/cuda-oxide/intrinsics/probes/sync_threads.ll index ce46214339..90ca7a294f 100644 --- a/cuda-oxide/intrinsics/probes/sync_threads.ll +++ b/cuda-oxide/intrinsics/probes/sync_threads.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tanh_approx_f32.ll b/cuda-oxide/intrinsics/probes/tanh_approx_f32.ll index d4fe0e6454..4214786041 100644 --- a/cuda-oxide/intrinsics/probes/tanh_approx_f32.ll +++ b/cuda-oxide/intrinsics/probes/tanh_approx_f32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_alloc.ll b/cuda-oxide/intrinsics/probes/tcgen05_alloc.ll index c4a98107a3..78eb03c2a9 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_alloc.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_alloc.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_alloc_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_alloc_cg2.ll index aa65343e80..ec3bf0049d 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_alloc_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_alloc_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_commit.ll b/cuda-oxide/intrinsics/probes/tcgen05_commit.ll index bd3f8dd917..f51ead80e0 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_commit.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_commit.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_commit_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_commit_cg2.ll index fd7ad4b836..887bd48078 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_commit_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_commit_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_commit_multicast.ll b/cuda-oxide/intrinsics/probes/tcgen05_commit_multicast.ll index 94c1b2c576..9e2498d08c 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_commit_multicast.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_commit_multicast.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_commit_multicast_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_commit_multicast_cg2.ll index f30b6e4148..8cbdbf4106 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_commit_multicast_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_commit_multicast_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_commit_shared_cluster.ll b/cuda-oxide/intrinsics/probes/tcgen05_commit_shared_cluster.ll index 80111f0d7f..43137a5b5f 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_commit_shared_cluster.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_commit_shared_cluster.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_commit_shared_cluster_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_commit_shared_cluster_cg2.ll index 1f1c0626d1..056111dec4 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_commit_shared_cluster_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_commit_shared_cluster_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b.ll index 82d6cb8fc2..52ac59336b 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b4x16_p64.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b4x16_p64.ll index befebd6d53..9f41f4a42f 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b4x16_p64.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b4x16_p64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b4x16_p64_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b4x16_p64_cg2.ll index d1e95e933a..f24277b543 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b4x16_p64_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b4x16_p64_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b6x16_p32.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b6x16_p32.ll index a005900b2d..ca1670f88b 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b6x16_p32.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b6x16_p32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b6x16_p32_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b6x16_p32_cg2.ll index 04732a8668..80283fe817 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b6x16_p32_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_b6x16_p32_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_cg2.ll index dc9fb404e8..1c61775cfd 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x128b_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b4x16_p64.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b4x16_p64.ll index c8e062f335..f9188397e7 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b4x16_p64.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b4x16_p64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b4x16_p64_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b4x16_p64_cg2.ll index af63d24990..ce7415b27a 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b4x16_p64_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b4x16_p64_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b6x16_p32.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b6x16_p32.ll index 7280645df7..9ab38c126c 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b6x16_p32.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b6x16_p32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b6x16_p32_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b6x16_p32_cg2.ll index eebc906521..0684f36216 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b6x16_p32_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_128x256b_b6x16_p32_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4.ll index 37461225c4..414a0fca71 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b4x16_p64.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b4x16_p64.ll index c93af18b1c..f4d985a3e9 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b4x16_p64.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b4x16_p64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b4x16_p64_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b4x16_p64_cg2.ll index d0f3edc2cc..c570845432 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b4x16_p64_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b4x16_p64_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b6x16_p32.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b6x16_p32.ll index 325ca3049f..af4accbcaa 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b6x16_p32.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b6x16_p32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b6x16_p32_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b6x16_p32_cg2.ll index 05c27fbd7d..480cdce7b9 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b6x16_p32_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_b6x16_p32_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_cg2.ll index 16562e1b5b..67b3e5e394 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_32x128b_warpx4_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b.ll index f56911c84f..fad35cb1e6 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b4x16_p64.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b4x16_p64.ll index 35301783a5..0ee873538f 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b4x16_p64.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b4x16_p64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b4x16_p64_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b4x16_p64_cg2.ll index 90c5f0f6cc..652c6a6e81 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b4x16_p64_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b4x16_p64_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b6x16_p32.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b6x16_p32.ll index bca748c663..3d5e40bc11 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b6x16_p32.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b6x16_p32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b6x16_p32_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b6x16_p32_cg2.ll index 702833f89d..0c05400522 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b6x16_p32_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_b6x16_p32_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_cg2.ll index c15a0380a9..4f4bc6989f 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_4x256b_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23.ll index 9ba6a1afb0..278e650628 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b4x16_p64.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b4x16_p64.ll index da556f5c37..a80cfc2424 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b4x16_p64.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b4x16_p64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b4x16_p64_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b4x16_p64_cg2.ll index 0f509be737..00f8d805e6 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b4x16_p64_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b4x16_p64_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b6x16_p32.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b6x16_p32.ll index 4bfc1eb567..4021f50cf6 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b6x16_p32.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b6x16_p32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b6x16_p32_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b6x16_p32_cg2.ll index ebbc1091e1..729bd592b7 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b6x16_p32_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_b6x16_p32_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_cg2.ll index c550717d81..451c111368 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_01_23_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13.ll index 63676f4103..3641306908 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b4x16_p64.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b4x16_p64.ll index 58031ab53e..ea45720aa8 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b4x16_p64.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b4x16_p64.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b4x16_p64_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b4x16_p64_cg2.ll index dca3dc3f61..078bb330a7 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b4x16_p64_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b4x16_p64_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b6x16_p32.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b6x16_p32.ll index 8f56444bde..b9b54670bb 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b6x16_p32.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b6x16_p32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b6x16_p32_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b6x16_p32_cg2.ll index 0e39bc6bb1..1df47e63d3 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b6x16_p32_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_b6x16_p32_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_cg2.ll index 4ad10e1aaa..e1412749ab 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_64x128b_warpx2_02_13_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_smem_to_tmem.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_smem_to_tmem.ll index 4c6c331e69..994ef9b21e 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_smem_to_tmem.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_smem_to_tmem.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_cp_smem_to_tmem_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_cp_smem_to_tmem_cg2.ll index 831794afa1..d499f8d012 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_cp_smem_to_tmem_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_cp_smem_to_tmem_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_dealloc.ll b/cuda-oxide/intrinsics/probes/tcgen05_dealloc.ll index 4b12b5229c..5793c9e4bd 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_dealloc.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_dealloc.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_dealloc_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_dealloc_cg2.ll index a2290352bb..c6b12d41f3 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_dealloc_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_dealloc_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_fence_after_thread_sync.ll b/cuda-oxide/intrinsics/probes/tcgen05_fence_after_thread_sync.ll index 0de274f6d8..ad5878d884 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_fence_after_thread_sync.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_fence_after_thread_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_fence_before_thread_sync.ll b/cuda-oxide/intrinsics/probes/tcgen05_fence_before_thread_sync.ll index 6c3adf446d..58a0f1a15e 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_fence_before_thread_sync.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_fence_before_thread_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x16_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x16_pack16.ll index 4cd1939041..943ae6853c 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x16_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x16_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x16_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x16_raw.ll index df340aff19..441955ce2c 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x16_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x16_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x1_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x1_pack16.ll index ec2e72598a..a195cade58 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x1_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x1_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x1_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x1_raw.ll index 1fe35c74da..94fa320fa3 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x1_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x1_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x2_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x2_pack16.ll index f4bb6baa5b..b625851f40 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x2_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x2_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x2_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x2_raw.ll index fc3b7238fe..ee582d5c14 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x2_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x2_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x32_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x32_pack16.ll index b1ec497d34..4abc2ccb2d 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x32_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x32_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x32_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x32_raw.ll index 8eb142e021..ebc3ce594f 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x32_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x32_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x4_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x4_pack16.ll index 02ad611b67..f950eea859 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x4_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x4_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x4_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x4_raw.ll index 1c150b2976..f589b00b03 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x4_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x4_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x64_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x64_pack16.ll index 3e03b7bf69..0b481302d7 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x64_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x64_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x64_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x64_raw.ll index 1608c66fed..13823a35c0 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x64_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x64_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x8_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x8_pack16.ll index 8534fbda34..e4463e7a79 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x8_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x8_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x8_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x8_raw.ll index f76772489a..b31b8892c7 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x8_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x128b_x8_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_pure.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_pure.ll index 7fae8fb55c..5d39b18dd1 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_pure.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_pure.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x16_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x16_pack16.ll index 0105e9e8ea..8f3b785428 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x16_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x16_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x16_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x16_raw.ll index 958d3fc0a5..8b844226bd 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x16_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x16_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x1_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x1_pack16.ll index 7502cfb83f..2c53dfe751 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x1_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x1_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x1_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x1_raw.ll index 68bb0696d2..0ac7ef4d19 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x1_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x1_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x2_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x2_pack16.ll index 9e57fd194d..d2f013db50 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x2_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x2_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x2_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x2_raw.ll index 872845b67a..55e38f34ba 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x2_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x2_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x32_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x32_pack16.ll index 5afb3694b2..1c53be801d 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x32_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x32_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x32_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x32_raw.ll index c028128445..b34689036b 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x32_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x32_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x4_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x4_pack16.ll index c8f706e70a..5947255f6e 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x4_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x4_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x4_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x4_raw.ll index 1b952c7b7b..fb374205b5 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x4_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x4_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x8_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x8_pack16.ll index 06335338a1..7ff7ccbd50 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x8_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x8_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x8_pure.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x8_pure.ll index 3dcd3ade49..622972f100 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x8_pure.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x8_pure.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x8_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x8_raw.ll index f768ce04e1..9e01612da3 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x8_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x256b_x8_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x128_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x128_pack16.ll index e59c0edd99..cce8d03411 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x128_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x128_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x128_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x128_raw.ll index 47e747bf5d..562da3d5fe 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x128_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x128_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x16_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x16_pack16.ll index 8a8f6105c7..daf6371a44 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x16_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x16_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x16_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x16_raw.ll index f008604758..d6ba9c877a 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x16_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x16_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x1_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x1_pack16.ll index 86f8f0cbb0..bdc36df7ca 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x1_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x1_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x1_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x1_raw.ll index 423665f378..7e9849c845 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x1_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x1_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x2_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x2_pack16.ll index 4127c7471e..c52672ea3f 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x2_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x2_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x2_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x2_raw.ll index d418c5d574..0a644f5e3f 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x2_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x2_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x32_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x32_pack16.ll index 76fe64abfe..a361cbe9e5 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x32_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x32_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x32_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x32_raw.ll index 6279678a5f..07c3396b74 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x32_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x32_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x4_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x4_pack16.ll index 5146366f90..c1e9337245 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x4_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x4_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x4_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x4_raw.ll index fd6f4846f6..363134901a 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x4_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x4_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x64_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x64_pack16.ll index ead837807b..93e18e2a92 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x64_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x64_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x64_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x64_raw.ll index 844a4c51cb..14d5946742 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x64_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x64_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x8_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x8_pack16.ll index 49aef2f670..b6a44f099e 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x8_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x8_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x8_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x8_raw.ll index 78701b3a4b..3c151dd431 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x8_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x32bx2_x8_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x128_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x128_pack16.ll index 5274d630cd..f43f5725fd 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x128_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x128_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x128_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x128_raw.ll index 3eff415d00..c70832daa8 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x128_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x128_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x16_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x16_pack16.ll index 78242360dc..b07b9b6963 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x16_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x16_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x16_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x16_raw.ll index b9a0a6defa..2477ac3eab 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x16_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x16_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x1_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x1_pack16.ll index b98a6710a6..fcae4f43b4 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x1_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x1_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x1_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x1_raw.ll index 19daa07248..2f59bf5e91 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x1_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x1_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x2_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x2_pack16.ll index 090fa7f370..afa3c367d9 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x2_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x2_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x2_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x2_raw.ll index 775140988f..3399e11252 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x2_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x2_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x32_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x32_pack16.ll index 8a217784e7..42eca7eb98 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x32_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x32_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x32_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x32_raw.ll index b75adec2ac..1925b6dada 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x32_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x32_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x4_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x4_pack16.ll index d2f9d27b1f..4614dcbded 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x4_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x4_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x4_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x4_raw.ll index 388cadf957..78c4def494 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x4_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x4_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x64_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x64_pack16.ll index b39a98a6a4..26451019f4 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x64_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x64_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x64_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x64_raw.ll index 793465df1b..46f93b5852 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x64_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x64_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x8_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x8_pack16.ll index 9a5c8ff5ed..fcc042c2bc 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x8_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x8_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x8_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x8_raw.ll index 3b8d845e34..4e19ab4f74 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x8_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_16x64b_x8_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x128_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x128_pack16.ll index f2bdcd3112..4ff8ed19a0 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x128_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x128_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x128_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x128_raw.ll index caf5a18b2b..3cc175e1b2 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x128_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x128_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x16_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x16_pack16.ll index dc66b087cd..ff9bcea403 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x16_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x16_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x16_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x16_raw.ll index 3d8c42c6a5..2de65f9f22 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x16_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x16_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x1_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x1_pack16.ll index c3fd537a04..b05b0d9903 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x1_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x1_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x1_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x1_raw.ll index a932d66bf8..c3ce95cda9 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x1_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x1_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x2_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x2_pack16.ll index 8d638f44c9..56d1a6352d 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x2_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x2_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x2_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x2_raw.ll index 84fa9cd7ac..e38232ec80 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x2_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x2_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x32_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x32_pack16.ll index 4fe43f6019..c8311df293 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x32_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x32_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x32_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x32_raw.ll index 15be748227..5548b1d311 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x32_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x32_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x4_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x4_pack16.ll index 1f4c7d2879..b39908bb10 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x4_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x4_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x4_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x4_raw.ll index 0546d73a29..7cefd0417b 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x4_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x4_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x64_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x64_pack16.ll index e9a8e5b826..ce7a14b48a 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x64_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x64_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x64_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x64_raw.ll index fd11c27900..a858ac930b 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x64_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x64_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x8_pack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x8_pack16.ll index aaeb633bda..2379d1e201 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x8_pack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x8_pack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x8_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x8_raw.ll index f25ac5f59a..eeb45a2270 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x8_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_ld_32x32b_x8_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_load_wait.ll b/cuda-oxide/intrinsics/probes/tcgen05_load_wait.ll index 833bce51c7..bb6a0ab590 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_load_wait.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_load_wait.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_e2m1.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_e2m1.ll index 736d26d5ec..7736f83e05 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_e2m3.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_e2m3.ll index c378b7628f..66c12f06f2 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_e3m2.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_e3m2.ll index a8e2b80df3..97eebba1a2 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_e4m3.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_e4m3.ll index 2fc78f0d0c..60fa5dfc28 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_e5m2.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_e5m2.ll index 2495ba13b5..146674e0ce 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_f16.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_f16.ll index d15ff6fb61..02932cfac5 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_f16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_f16_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_f16_cg2.ll index 1820c4b7b4..38869213f6 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_f16_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_f16_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_shared.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_shared.ll index ce36c8ec2c..ff3d9e79ab 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_shared.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_shared.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_sp_shared.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_sp_shared.ll index e39ec244cf..cbe63e2f92 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_sp_shared.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_sp_shared.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_sp_tensor.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_sp_tensor.ll index 10b8427729..ac0f3627ff 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_sp_tensor.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_sp_tensor.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_sp_tensor_ashift.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_sp_tensor_ashift.ll index c07c86c7ae..5f94aa70e5 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_sp_tensor_ashift.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_sp_tensor_ashift.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_tensor.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_tensor.ll index f72ee49cac..f72e347418 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_tensor.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_tensor.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_tensor_ashift.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_tensor_ashift.ll index 39a1cc6552..d72a956ec6 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_tensor_ashift.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_tensor_ashift.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_bf16.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_bf16.ll index a3d1fd74a5..53ad38544e 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_bf16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_bf16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e2m1.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e2m1.ll index 4b2bb03687..8270acf81c 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e2m1.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e2m1.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e2m3.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e2m3.ll index 7cf23ae50f..c3d2b9f01c 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e2m3.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e2m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e3m2.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e3m2.ll index adfe7e9dde..9e615f9939 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e3m2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e3m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e4m3.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e4m3.ll index 0caa820372..470a7ec0ce 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e4m3.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e4m3.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e5m2.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e5m2.ll index f373132009..c14ae967ce 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e5m2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_e5m2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_f16.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_f16.ll index 62eac4d1bb..e3005220aa 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_f16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_f16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_shared.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_shared.ll index 10de4733d1..af582b2b6c 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_shared.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_shared.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_shared_zero_col_mask.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_shared_zero_col_mask.ll index ae0b0bce13..83114a4425 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_shared_zero_col_mask.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_shared_zero_col_mask.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_shared.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_shared.ll index 769ffbb069..a63ab88da6 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_shared.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_shared.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_shared_zero_col_mask.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_shared_zero_col_mask.ll index 5257356c11..f9a6974faa 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_shared_zero_col_mask.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_shared_zero_col_mask.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_tensor.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_tensor.ll index 8d5f3e3467..f269546d0a 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_tensor.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_tensor.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_tensor_zero_col_mask.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_tensor_zero_col_mask.ll index 204782ff2e..a54bfcb3ab 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_tensor_zero_col_mask.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_sp_tensor_zero_col_mask.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_tensor.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_tensor.ll index eb6e5634ac..9cc7426bf2 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_tensor.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_tensor.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_tensor_zero_col_mask.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_tensor_zero_col_mask.ll index 10f8b0710a..4405d54481 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_tensor_zero_col_mask.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_tensor_zero_col_mask.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_tf32.ll b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_tf32.ll index 45f263ce41..dd5a638957 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_tf32.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_mma_ws_tf32.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_relinquish_alloc_permit.ll b/cuda-oxide/intrinsics/probes/tcgen05_relinquish_alloc_permit.ll index 36a3685001..ea90485894 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_relinquish_alloc_permit.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_relinquish_alloc_permit.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_relinquish_alloc_permit_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_relinquish_alloc_permit_cg2.ll index d4266950d4..38750abff9 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_relinquish_alloc_permit_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_relinquish_alloc_permit_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_shift_down.ll b/cuda-oxide/intrinsics/probes/tcgen05_shift_down.ll index 3ccc0f975f..efd751ea26 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_shift_down.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_shift_down.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_shift_down_cg2.ll b/cuda-oxide/intrinsics/probes/tcgen05_shift_down_cg2.ll index db0ecaca04..682660c16e 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_shift_down_cg2.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_shift_down_cg2.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x16_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x16_raw.ll index 9c2e8989ea..6bd5e18a1f 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x16_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x16_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x16_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x16_unpack16.ll index d4c46aeb62..da2f62e2ef 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x16_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x16_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x1_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x1_raw.ll index 5edf5b5071..baf53e945b 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x1_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x1_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x1_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x1_unpack16.ll index d01f0d50ea..160ebe8304 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x1_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x1_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x2_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x2_raw.ll index 6e846371cb..b07fb575f8 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x2_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x2_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x2_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x2_unpack16.ll index 52cf3259c8..c57a143fd6 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x2_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x2_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x32_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x32_raw.ll index 90c75f2302..15f94fe28a 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x32_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x32_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x32_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x32_unpack16.ll index fd6480b738..4479311fe3 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x32_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x32_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x4_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x4_raw.ll index 509bc3cb19..d1caadb1a0 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x4_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x4_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x4_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x4_unpack16.ll index 922db84df1..b2c0ac5d66 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x4_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x4_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x64_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x64_raw.ll index 2eb469bc29..36c9b30938 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x64_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x64_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x64_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x64_unpack16.ll index f8235fac98..7af2de5062 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x64_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x64_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x8_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x8_raw.ll index 5105d6755a..4f932e6605 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x8_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x8_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x8_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x8_unpack16.ll index 091efa823b..f2e160b6a5 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x8_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x128b_x8_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x16_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x16_raw.ll index 6744944547..210c9f8f4d 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x16_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x16_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x16_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x16_unpack16.ll index 654c677c8e..29e2242e60 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x16_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x16_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x1_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x1_raw.ll index 4c9dedc9a2..fda93e447a 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x1_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x1_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x1_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x1_unpack16.ll index 38e23fa423..ac2e9ce2f6 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x1_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x1_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x2_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x2_raw.ll index efa9ff5bc8..5be863aae3 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x2_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x2_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x2_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x2_unpack16.ll index 5214de84be..0ff1d22f5d 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x2_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x2_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x32_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x32_raw.ll index 6f094eb6c9..c892ca29b4 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x32_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x32_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x32_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x32_unpack16.ll index 4d533bcfaf..80211d7e1e 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x32_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x32_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x4_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x4_raw.ll index 10e1fdb8ba..31a9282ab0 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x4_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x4_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x4_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x4_unpack16.ll index 0ca4a5e078..740ab8dcf8 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x4_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x4_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x8_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x8_raw.ll index b1bb736b7c..2a364d2472 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x8_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x8_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x8_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x8_unpack16.ll index 80f2d438c7..19e7bbad29 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x8_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x256b_x8_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x128_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x128_raw.ll index 88c82bea1e..e921053527 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x128_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x128_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x128_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x128_unpack16.ll index c983949996..17b38b5f34 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x128_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x128_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x16_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x16_raw.ll index baf88fd567..ca727f5c70 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x16_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x16_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x16_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x16_unpack16.ll index ec54132ec6..0ef4968320 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x16_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x16_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x1_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x1_raw.ll index 04b9c20add..4aee0bcd42 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x1_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x1_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x1_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x1_unpack16.ll index 25e96aed22..bb2a53b35b 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x1_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x1_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x2_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x2_raw.ll index c1e44ffcb3..6d284ea79c 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x2_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x2_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x2_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x2_unpack16.ll index 9a987bf9f4..177a3f5725 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x2_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x2_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x32_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x32_raw.ll index 8f5c45e08f..4aba04957e 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x32_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x32_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x32_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x32_unpack16.ll index 0be2af63f4..e2a8a82a0a 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x32_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x32_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x4_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x4_raw.ll index 75c89fde01..05eee361ad 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x4_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x4_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x4_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x4_unpack16.ll index f20c4daffa..bd1f2a35eb 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x4_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x4_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x64_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x64_raw.ll index 6ff9c8c77f..dc9a7f4785 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x64_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x64_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x64_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x64_unpack16.ll index db6fa748b7..1d16b9ab5c 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x64_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x64_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x8_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x8_raw.ll index 9e077f10ae..cb4bdb5ca1 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x8_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x8_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x8_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x8_unpack16.ll index 7bd0b29049..6437222b2e 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x8_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x32bx2_x8_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x128_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x128_raw.ll index e2c3c4cc4a..039aa50223 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x128_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x128_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x128_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x128_unpack16.ll index 06d695f163..05fbb6326f 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x128_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x128_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x16_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x16_raw.ll index ca6a11c0b0..4dd98f9106 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x16_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x16_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x16_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x16_unpack16.ll index 7f780048de..d3330e92a4 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x16_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x16_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x1_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x1_raw.ll index a3ea1cc760..dec8ce54eb 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x1_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x1_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x1_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x1_unpack16.ll index 84f1c666cd..0174ccb84d 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x1_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x1_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x2_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x2_raw.ll index 4f5b191e2c..3b2020be07 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x2_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x2_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x2_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x2_unpack16.ll index 434781b5e7..cfcd5a4ecf 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x2_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x2_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x32_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x32_raw.ll index ca703f9860..491ef09a52 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x32_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x32_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x32_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x32_unpack16.ll index 494700de28..4e4b0254fd 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x32_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x32_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x4_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x4_raw.ll index c924d7ae36..a5843f37cb 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x4_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x4_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x4_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x4_unpack16.ll index f20e110cd3..e2fd28ee8d 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x4_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x4_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x64_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x64_raw.ll index 69f69eb029..2e807190c3 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x64_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x64_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x64_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x64_unpack16.ll index 8a63ab3478..cc8095320e 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x64_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x64_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x8_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x8_raw.ll index 3b63f507fd..6f6c421fc9 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x8_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x8_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x8_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x8_unpack16.ll index 1c4bec7c72..832ee97334 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x8_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_16x64b_x8_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x128_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x128_raw.ll index c6a6911364..b1f265dc54 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x128_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x128_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x128_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x128_unpack16.ll index 6485605e9a..4372203efa 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x128_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x128_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x16_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x16_raw.ll index 86b42eee8e..da4a59a3b7 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x16_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x16_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x16_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x16_unpack16.ll index dd834d07d0..29ae893dee 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x16_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x16_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x1_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x1_raw.ll index 51cedb3e24..eddcd92c8b 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x1_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x1_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x1_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x1_unpack16.ll index 145815e2da..b696b44b58 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x1_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x1_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x2_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x2_raw.ll index e9e6cc6ca2..171613de3a 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x2_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x2_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x2_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x2_unpack16.ll index d0c0c539b5..af55c50077 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x2_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x2_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x32_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x32_raw.ll index f9da8f8403..de783aa54b 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x32_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x32_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x32_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x32_unpack16.ll index fe613f7fcb..87aa39d1b6 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x32_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x32_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x4_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x4_raw.ll index 416c390ca4..c8875ba80b 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x4_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x4_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x4_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x4_unpack16.ll index 02de245621..a420186f9c 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x4_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x4_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x64_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x64_raw.ll index a7866b2cd4..a92cde652c 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x64_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x64_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x64_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x64_unpack16.ll index d40f1f8bf8..6695ccc474 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x64_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x64_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x8_raw.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x8_raw.ll index 3663152919..fe841b7120 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x8_raw.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x8_raw.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x8_unpack16.ll b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x8_unpack16.ll index e40810dbc0..7656d3f7e5 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x8_unpack16.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_st_32x32b_x8_unpack16.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tcgen05_store_wait.ll b/cuda-oxide/intrinsics/probes/tcgen05_store_wait.ll index 4a1646c6ae..91f6b243d1 100644 --- a/cuda-oxide/intrinsics/probes/tcgen05_store_wait.ll +++ b/cuda-oxide/intrinsics/probes/tcgen05_store_wait.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tensormap_replace_box_dim.ll b/cuda-oxide/intrinsics/probes/tensormap_replace_box_dim.ll index 7d5ad074f6..335991879b 100644 --- a/cuda-oxide/intrinsics/probes/tensormap_replace_box_dim.ll +++ b/cuda-oxide/intrinsics/probes/tensormap_replace_box_dim.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tensormap_replace_element_stride.ll b/cuda-oxide/intrinsics/probes/tensormap_replace_element_stride.ll index c9620b9094..90a9ff6954 100644 --- a/cuda-oxide/intrinsics/probes/tensormap_replace_element_stride.ll +++ b/cuda-oxide/intrinsics/probes/tensormap_replace_element_stride.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tensormap_replace_element_type.ll b/cuda-oxide/intrinsics/probes/tensormap_replace_element_type.ll index ef6884fc13..228ed523bb 100644 --- a/cuda-oxide/intrinsics/probes/tensormap_replace_element_type.ll +++ b/cuda-oxide/intrinsics/probes/tensormap_replace_element_type.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tensormap_replace_fill_mode.ll b/cuda-oxide/intrinsics/probes/tensormap_replace_fill_mode.ll index a817b06580..77bbdfb025 100644 --- a/cuda-oxide/intrinsics/probes/tensormap_replace_fill_mode.ll +++ b/cuda-oxide/intrinsics/probes/tensormap_replace_fill_mode.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tensormap_replace_global_address.ll b/cuda-oxide/intrinsics/probes/tensormap_replace_global_address.ll index aa1c74e519..97379ffe52 100644 --- a/cuda-oxide/intrinsics/probes/tensormap_replace_global_address.ll +++ b/cuda-oxide/intrinsics/probes/tensormap_replace_global_address.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tensormap_replace_global_dim.ll b/cuda-oxide/intrinsics/probes/tensormap_replace_global_dim.ll index 7e169639ec..89fb8819e6 100644 --- a/cuda-oxide/intrinsics/probes/tensormap_replace_global_dim.ll +++ b/cuda-oxide/intrinsics/probes/tensormap_replace_global_dim.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tensormap_replace_global_stride.ll b/cuda-oxide/intrinsics/probes/tensormap_replace_global_stride.ll index 03eaa106f4..cb05170cee 100644 --- a/cuda-oxide/intrinsics/probes/tensormap_replace_global_stride.ll +++ b/cuda-oxide/intrinsics/probes/tensormap_replace_global_stride.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tensormap_replace_interleave_layout.ll b/cuda-oxide/intrinsics/probes/tensormap_replace_interleave_layout.ll index de6506df8a..1d6ba891a1 100644 --- a/cuda-oxide/intrinsics/probes/tensormap_replace_interleave_layout.ll +++ b/cuda-oxide/intrinsics/probes/tensormap_replace_interleave_layout.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tensormap_replace_rank.ll b/cuda-oxide/intrinsics/probes/tensormap_replace_rank.ll index fe71f64549..d77b27382f 100644 --- a/cuda-oxide/intrinsics/probes/tensormap_replace_rank.ll +++ b/cuda-oxide/intrinsics/probes/tensormap_replace_rank.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tensormap_replace_swizzle_atomicity.ll b/cuda-oxide/intrinsics/probes/tensormap_replace_swizzle_atomicity.ll index ad64dd0bc9..c23ad58e62 100644 --- a/cuda-oxide/intrinsics/probes/tensormap_replace_swizzle_atomicity.ll +++ b/cuda-oxide/intrinsics/probes/tensormap_replace_swizzle_atomicity.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/tensormap_replace_swizzle_mode.ll b/cuda-oxide/intrinsics/probes/tensormap_replace_swizzle_mode.ll index fc789edcb3..f26d3c06e7 100644 --- a/cuda-oxide/intrinsics/probes/tensormap_replace_swizzle_mode.ll +++ b/cuda-oxide/intrinsics/probes/tensormap_replace_swizzle_mode.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/thread_idx_x.ll b/cuda-oxide/intrinsics/probes/thread_idx_x.ll index 1b09367f16..d7383dc1a6 100644 --- a/cuda-oxide/intrinsics/probes/thread_idx_x.ll +++ b/cuda-oxide/intrinsics/probes/thread_idx_x.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/thread_idx_y.ll b/cuda-oxide/intrinsics/probes/thread_idx_y.ll index 365fe70767..9c524bc93e 100644 --- a/cuda-oxide/intrinsics/probes/thread_idx_y.ll +++ b/cuda-oxide/intrinsics/probes/thread_idx_y.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/thread_idx_z.ll b/cuda-oxide/intrinsics/probes/thread_idx_z.ll index e6800e32eb..99747f6487 100644 --- a/cuda-oxide/intrinsics/probes/thread_idx_z.ll +++ b/cuda-oxide/intrinsics/probes/thread_idx_z.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/threadfence.ll b/cuda-oxide/intrinsics/probes/threadfence.ll index 1e7570c082..5841905be8 100644 --- a/cuda-oxide/intrinsics/probes/threadfence.ll +++ b/cuda-oxide/intrinsics/probes/threadfence.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/threadfence_block.ll b/cuda-oxide/intrinsics/probes/threadfence_block.ll index d9df7a1120..6ebe32fea8 100644 --- a/cuda-oxide/intrinsics/probes/threadfence_block.ll +++ b/cuda-oxide/intrinsics/probes/threadfence_block.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/threadfence_system.ll b/cuda-oxide/intrinsics/probes/threadfence_system.ll index 8b8dd48126..147879f7f7 100644 --- a/cuda-oxide/intrinsics/probes/threadfence_system.ll +++ b/cuda-oxide/intrinsics/probes/threadfence_system.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/total_smem_size.lib_nvvm.ll b/cuda-oxide/intrinsics/probes/total_smem_size.lib_nvvm.ll index 1ec13d5690..b919617aef 100644 --- a/cuda-oxide/intrinsics/probes/total_smem_size.lib_nvvm.ll +++ b/cuda-oxide/intrinsics/probes/total_smem_size.lib_nvvm.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/total_smem_size.llvm_nvptx.ll b/cuda-oxide/intrinsics/probes/total_smem_size.llvm_nvptx.ll index ad81b5ba62..18ae64f1ee 100644 --- a/cuda-oxide/intrinsics/probes/total_smem_size.llvm_nvptx.ll +++ b/cuda-oxide/intrinsics/probes/total_smem_size.llvm_nvptx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/trap.ll b/cuda-oxide/intrinsics/probes/trap.ll index a20d4fd79c..53898f8de2 100644 --- a/cuda-oxide/intrinsics/probes/trap.ll +++ b/cuda-oxide/intrinsics/probes/trap.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/uni_sync.ll b/cuda-oxide/intrinsics/probes/uni_sync.ll index 0a2c1a8862..03fe54c6db 100644 --- a/cuda-oxide/intrinsics/probes/uni_sync.ll +++ b/cuda-oxide/intrinsics/probes/uni_sync.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/warpid.lib_nvvm.ll b/cuda-oxide/intrinsics/probes/warpid.lib_nvvm.ll index fc97dc0ef4..67b272cf96 100644 --- a/cuda-oxide/intrinsics/probes/warpid.lib_nvvm.ll +++ b/cuda-oxide/intrinsics/probes/warpid.lib_nvvm.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/warpid.llvm_nvptx.ll b/cuda-oxide/intrinsics/probes/warpid.llvm_nvptx.ll index d67ebdec7d..4ded8940bc 100644 --- a/cuda-oxide/intrinsics/probes/warpid.llvm_nvptx.ll +++ b/cuda-oxide/intrinsics/probes/warpid.llvm_nvptx.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/wgmma_commit_group.ll b/cuda-oxide/intrinsics/probes/wgmma_commit_group.ll index ae20244245..8d65c25f1b 100644 --- a/cuda-oxide/intrinsics/probes/wgmma_commit_group.ll +++ b/cuda-oxide/intrinsics/probes/wgmma_commit_group.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/wgmma_fence.ll b/cuda-oxide/intrinsics/probes/wgmma_fence.ll index 796daf4114..854ff0e9dd 100644 --- a/cuda-oxide/intrinsics/probes/wgmma_fence.ll +++ b/cuda-oxide/intrinsics/probes/wgmma_fence.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/intrinsics/probes/wgmma_wait_group.ll b/cuda-oxide/intrinsics/probes/wgmma_wait_group.ll index f1cb4efa55..ae3b6b7ec9 100644 --- a/cuda-oxide/intrinsics/probes/wgmma_wait_group.ll +++ b/cuda-oxide/intrinsics/probes/wgmma_wait_group.ll @@ -1,7 +1,7 @@ ; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ; SPDX-License-Identifier: Apache-2.0 ; @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 4b1945f6e57a1cd71e7dcf59ffdfba1c6aa7827e72a1fd21232d9327c350c719. +; Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: b8c889f1ecb3e7e4778d8347bdfa684c1664c72050fc20fb6dbe5d5f988cc81d. ; LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. target triple = "nvptx64-nvidia-cuda" diff --git a/cuda-oxide/scripts/check-spdx-headers.sh b/cuda-oxide/scripts/check-spdx-headers.sh index 13927665f1..ba6314e352 100755 --- a/cuda-oxide/scripts/check-spdx-headers.sh +++ b/cuda-oxide/scripts/check-spdx-headers.sh @@ -69,6 +69,11 @@ BSD_EXEMPT_FILES=( cuda-oxide/cuteir/examples/blockscale_gemm_cute/reference/dense_blockscaled_gemm_persistent_cooperative.py cuda-oxide/cuteir/examples/blockscale_gemm_cute/reference/blockscaled_gemm_dispatch.py cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/reference/flash_attention_v2.py + cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/cute/blackwell/kernel/attention/fmha/fmha.py + cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/helpers/fmha_helpers.py + cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/helpers/__init__.py + cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/cute/blackwell/kernel/attention/mla/mla_decode_fp16.py + cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/cute/blackwell/kernel/attention/mla/mla_helpers.py ) # `head | grep -q` under pipefail can report a match as failure when grep From d642dcbfbf8c2f249a6f07307c158c5c0e89a4f7 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Wed, 14 Oct 2026 11:30:59 +0530 Subject: [PATCH 34/99] Modularize and document the CuTe implementation Organize CuTe layouts, copies, MMA, pipelines, importer handlers, and MLIR exporters into focused modules. Keep public reexports stable and update exact compiler recognition paths with regression coverage. Route SM120a attention primitives through CTM, including kernel macros and storage reexports, and retain the migration's correctness, instruction, and performance evidence. Add concise Rust documentation for public, private, and restricted items, including compiler hooks, fields, macros, examples, and test helpers. Refresh diagnostic snapshots for documentation-related source locations and repair links to the reorganized implementation. Signed-off-by: nihalpasham --- .../crates/cuda-oxide-mlir-export/README.md | 31 +- .../cuda-oxide-mlir-export/src/cute/common.rs | 427 ++ .../src/cute/common_sm100.rs | 22 + .../src/cute/copy/cooperative.rs | 135 + .../src/cute/copy/emit.rs | 254 + .../src/cute/copy/mod.rs | 12 + .../src/cute/copy/registry.rs | 30 + .../src/cute/copy/sm100_attention.rs | 91 + .../src/cute/copy/tma.rs | 207 + .../cuda-oxide-mlir-export/src/cute/emit.rs | 272 + .../src/cute/epilogue/mod.rs | 9 + .../src/cute/epilogue/sm100.rs | 88 + .../src/cute/epilogue/sm120.rs | 547 ++ .../src/cute/independent_sm100.rs | 23 + .../src/cute/mma/ldmatrix.rs | 290 ++ .../src/cute/mma/mod.rs | 11 + .../src/cute/mma/sm100.rs | 65 + .../mma/sm100_attention.rs} | 108 +- .../src/cute/mma/sm120.rs | 1083 ++++ .../cuda-oxide-mlir-export/src/cute/mod.rs | 36 + .../src/cute/pipeline/mod.rs | 9 + .../src/cute/pipeline/sm100.rs | 87 + .../src/cute/pipeline/sm120.rs | 618 +++ .../src/cute/registry.rs | 174 + .../src/cute/scheduler/mod.rs | 8 + .../src/cute/scheduler/sm120.rs | 271 + .../src/{cute_sm100.rs => cute/sm100.rs} | 512 +- .../{cute.rs => cute/tensor/elementwise.rs} | 40 + .../src/{cute_gemv.rs => cute/tensor/gemv.rs} | 52 +- .../src/cute/tensor/mod.rs | 9 + .../tests/cooperative.rs} | 642 +-- .../src/cute/tests/gemm.rs | 128 + .../cuda-oxide-mlir-export/src/cute_gemm.rs | 3118 ----------- .../src/grid_constant.rs | 6 +- .../crates/cuda-oxide-mlir-export/src/lib.rs | 62 +- .../src/{mir_checked.rs => mir/checked.rs} | 46 +- .../src/{mir_core.rs => mir/core.rs} | 80 +- .../src/{mir_memory.rs => mir/memory.rs} | 56 +- .../cuda-oxide-mlir-export/src/mir/mod.rs | 11 + .../src/{mir_shared.rs => mir/shared.rs} | 9 +- .../{nvvm_attention.rs => nvvm/attention.rs} | 9 +- .../src/{nvvm_cluster.rs => nvvm/cluster.rs} | 100 +- .../cuda-oxide-mlir-export/src/nvvm/mod.rs | 14 + .../src/{nvvm_numeric.rs => nvvm/numeric.rs} | 41 +- .../src/{nvvm_sm120.rs => nvvm/sm120.rs} | 11 + .../src/{nvvm_sregs.rs => nvvm/sregs.rs} | 12 +- .../sync_restrict.rs} | 5 + .../src/{nvvm_tcgen05.rs => nvvm/tcgen05.rs} | 46 +- .../cuda-oxide-mlir-export/src/packs.rs | 18 + .../cuda-oxide-mlir-export/src/primitives.rs | 15 +- .../cuda-oxide-mlir-export/src/profile.rs | 56 +- .../stack-temporaries.md | 35 - .../tests/address_contract.rs | 5 + .../tests/attention_cute.rs | 3 + .../tests/attention_descriptor.rs | 9 + .../tests/attention_primitives.rs | 2 + .../tests/checked_storage.rs | 17 + .../tests/cute_elementwise.rs | 5 + .../cuda-oxide-mlir-export/tests/cute_gemv.rs | 7 + .../tests/grid_constant.rs | 5 + .../tests/nvvm_attention.rs | 23 + .../tests/primitives_profile.rs | 9 + .../tests/rust_float_max.rs | 8 + .../tests/sm120_primitives.rs | 26 + .../tests/sync_restrict.rs | 13 + .../emit.rs} | 13 +- .../{block_scaled.rs => block_scaled/mod.rs} | 35 +- .../src/translator/cute/classify.rs | 612 +++ .../src/translator/cute/dispatch.rs | 751 +++ .../mir-importer/src/translator/cute/emit.rs | 10 +- .../{epilogue_emit.rs => epilogue/emit.rs} | 21 +- .../cute/{epilogue.rs => epilogue/mod.rs} | 20 +- .../src/translator/cute/identity.rs | 374 ++ .../mir-importer/src/translator/cute/mod.rs | 2291 +-------- .../{pipeline_emit.rs => pipeline/emit.rs} | 36 +- .../cute/{pipeline.rs => pipeline/mod.rs} | 16 + .../src/translator/cute/registry.rs | 662 +++ .../{scheduler_emit.rs => scheduler/emit.rs} | 15 +- .../cute/{scheduler.rs => scheduler/mod.rs} | 12 + .../cute/{sm100_emit.rs => sm100.rs} | 31 +- ...0_attention_emit.rs => sm100_attention.rs} | 14 +- .../{smem_mma_emit.rs => smem_mma/emit.rs} | 50 +- .../cute/{smem_mma.rs => smem_mma/mod.rs} | 56 +- .../src/translator/cute/static_config.rs | 42 +- .../cute/{tensor_emit.rs => tensor/emit.rs} | 17 +- .../cute/{tensor.rs => tensor/mod.rs} | 20 + .../src/translator/rvalue/const_alloc.rs | 4 +- .../src/translator/rvalue/const_bytes.rs | 4 +- .../src/translator/rvalue/place_addr.rs | 4 +- .../rustc-codegen-cuda/src/collector.rs | 83 +- .../primitives_api_compile/src/kernel.rs | 42 +- cuda-oxide/cuteir/README.md | 39 + cuda-oxide/cuteir/ctm-rs/src/async_copy.rs | 38 + cuda-oxide/cuteir/ctm-rs/src/cluster.rs | 7 + cuda-oxide/cuteir/ctm-rs/src/convert.rs | 17 + cuda-oxide/cuteir/ctm-rs/src/float.rs | 29 + cuda-oxide/cuteir/ctm-rs/src/groups.rs | 5 + cuda-oxide/cuteir/ctm-rs/src/host/tma.rs | 32 +- cuda-oxide/cuteir/ctm-rs/src/lib.rs | 21 + cuda-oxide/cuteir/ctm-rs/src/memory.rs | 50 +- cuda-oxide/cuteir/ctm-rs/src/mma.rs | 49 +- .../cuteir/ctm-rs/src/tcgen05/descriptors.rs | 54 + cuda-oxide/cuteir/ctm-rs/src/tcgen05/mma.rs | 15 +- cuda-oxide/cuteir/ctm-rs/src/tcgen05/mod.rs | 7 + cuda-oxide/cuteir/ctm-rs/src/tcgen05/tmem.rs | 48 +- cuda-oxide/cuteir/ctm-rs/src/tma.rs | 14 + cuda-oxide/cuteir/ctm-rs/src/warp.rs | 13 +- .../cuteir/ctm-rs/tests/accumulator_tile.rs | 5 + .../cuteir/ctm-rs/tests/attention_api.rs | 11 + .../primitives_accumulator_tile_bounds.rs | 6 + .../primitives_accumulator_tile_roles.rs | 3 + .../primitives_accumulator_tile_roles.stderr | 20 +- .../primitives_accumulator_tile_rows.rs | 3 + .../primitives_accumulator_tile_rows.stderr | 12 +- .../compile_fail/primitives_address_domain.rs | 4 + .../primitives_address_domain.stderr | 26 +- .../primitives_async_wait_runtime.rs | 16 + .../primitives_async_wait_runtime.stderr | 7 + .../primitives_bulk_wait_count.rs | 3 + .../primitives_bulk_wait_count.stderr | 4 +- .../primitives_bulk_wait_runtime.rs | 3 + .../primitives_bulk_wait_runtime.stderr | 8 +- .../compile_fail/primitives_f16_mma_roles.rs | 16 + .../primitives_f16_mma_roles.stderr | 18 + .../compile_fail/primitives_guard_consumed.rs | 3 + .../primitives_guard_consumed.stderr | 4 +- .../primitives_invalid_columns.rs | 3 + .../primitives_invalid_columns.stderr | 4 +- .../compile_fail/primitives_mma_b_shape.rs | 5 + .../primitives_mma_b_shape.stderr | 4 +- .../primitives_mma_configuration.rs | 6 + .../primitives_mma_configuration.stderr | 24 +- .../compile_fail/primitives_mma_descriptor.rs | 5 + .../primitives_mma_descriptor.stderr | 4 +- .../primitives_mma_destination.rs | 5 + .../primitives_mma_destination.stderr | 4 +- .../compile_fail/primitives_mma_k_step.rs | 5 + .../compile_fail/primitives_mma_k_step.stderr | 48 +- .../primitives_mma_output_format.rs | 5 + .../primitives_mma_output_format.stderr | 64 +- .../compile_fail/primitives_mma_roles.rs | 5 + .../compile_fail/primitives_mma_roles.stderr | 12 +- .../primitives_no_implicit_allocation.rs | 4 + .../primitives_no_implicit_allocation.stderr | 12 +- .../compile_fail/primitives_remote_domain.rs | 3 + .../primitives_remote_domain.stderr | 4 +- .../primitives_remote_relaxed_domain.rs | 3 + .../primitives_remote_relaxed_domain.stderr | 26 +- .../primitives_tensor_map_element.rs | 3 + .../primitives_tensor_map_element.stderr | 4 +- .../compile_fail/primitives_tma_format.rs | 3 + .../compile_fail/primitives_tma_format.stderr | 4 +- .../primitives_unready_barrier.rs | 3 + .../primitives_unready_barrier.stderr | 4 +- .../primitives_view_after_deallocation.rs | 4 + .../primitives_view_after_deallocation.stderr | 10 +- .../compile_pass/attention_primitives.rs | 50 + .../ctm-rs/tests/primitives_mxf4_storage.rs | 11 +- .../cuteir/ctm-rs/tests/primitives_safety.rs | 2 + .../cuteir/ctm-rs/tests/primitives_sm120.rs | 2 + .../cuteir/ctm-rs/tests/primitives_types.rs | 6 + cuda-oxide/cuteir/cute-layout/src/algebra.rs | 14 +- cuda-oxide/cuteir/cute-layout/src/composed.rs | 18 +- cuda-oxide/cuteir/cute-layout/src/lib.rs | 24 +- .../cuteir/cute-layout/src/predication.rs | 32 +- cuda-oxide/cuteir/cute-layout/src/swizzle.rs | 9 + .../cuteir/cute-layout/src/validation.rs | 1668 ------ .../cute-layout/src/validation/cooperative.rs | 490 ++ .../cute-layout/src/validation/error.rs | 266 + .../cute-layout/src/validation/ldmatrix.rs | 88 + .../cuteir/cute-layout/src/validation/mod.rs | 86 + .../cute-layout/src/validation/tests.rs | 331 ++ .../src/validation/thread_value.rs | 366 ++ .../cuteir/cute-layout/src/validation/tma.rs | 136 + .../cute-rs/examples/global_store_compile.rs | 1 + cuda-oxide/cuteir/cute-rs/src/block_scaled.rs | 860 ---- .../cute-rs/src/block_scaled/__compiler.rs | 132 + .../cuteir/cute-rs/src/block_scaled/gemv.rs | 150 + .../cute-rs/src/block_scaled/layouts.rs | 219 + .../cuteir/cute-rs/src/block_scaled/mod.rs | 80 + .../cuteir/cute-rs/src/block_scaled/scales.rs | 100 + .../cuteir/cute-rs/src/block_scaled/tensor.rs | 170 + .../cuteir/cute-rs/src/block_scaled/tests.rs | 122 + .../cuteir/cute-rs/src/block_scaled_mma.rs | 965 ---- .../src/block_scaled_mma/__compiler.rs | 214 + .../cute-rs/src/block_scaled_mma/epilogue.rs | 29 + .../cute-rs/src/block_scaled_mma/fragments.rs | 308 ++ .../cute-rs/src/block_scaled_mma/mod.rs | 59 + .../cute-rs/src/block_scaled_mma/tests.rs | 133 + .../cute-rs/src/block_scaled_mma/tiled.rs | 338 ++ cuda-oxide/cuteir/cute-rs/src/cooperative.rs | 2 + cuda-oxide/cuteir/cute-rs/src/epilogue.rs | 21 + cuda-oxide/cuteir/cute-rs/src/lib.rs | 26 +- cuda-oxide/cuteir/cute-rs/src/markers.rs | 18 +- cuda-oxide/cuteir/cute-rs/src/mma.rs | 23 +- cuda-oxide/cuteir/cute-rs/src/numeric.rs | 13 +- cuda-oxide/cuteir/cute-rs/src/pipeline.rs | 48 + cuda-oxide/cuteir/cute-rs/src/scheduler.rs | 28 + cuda-oxide/cuteir/cute-rs/src/sm100.rs | 522 -- .../cuteir/cute-rs/src/sm100/__compiler.rs | 137 + cuda-oxide/cuteir/cute-rs/src/sm100/mma.rs | 98 + cuda-oxide/cuteir/cute-rs/src/sm100/mod.rs | 39 + .../cuteir/cute-rs/src/sm100/pipeline.rs | 190 + cuda-oxide/cuteir/cute-rs/src/sm100/shared.rs | 34 + cuda-oxide/cuteir/cute-rs/src/sm100/tests.rs | 20 + cuda-oxide/cuteir/cute-rs/src/sm100/tma.rs | 49 + cuda-oxide/cuteir/cute-rs/src/sm100/tmem.rs | 94 + .../cuteir/cute-rs/src/sm100_attention.rs | 536 -- .../cute-rs/src/sm100_attention/__compiler.rs | 73 + .../cute-rs/src/sm100_attention/host.rs | 5 +- .../cuteir/cute-rs/src/sm100_attention/mma.rs | 130 + .../cuteir/cute-rs/src/sm100_attention/mod.rs | 28 + .../cute-rs/src/sm100_attention/pipeline.rs | 105 + .../cute-rs/src/sm100_attention/shared.rs | 168 + .../cuteir/cute-rs/src/sm100_attention/tma.rs | 60 + .../cute-rs/src/sm100_attention/tmem.rs | 189 + cuda-oxide/cuteir/cute-rs/src/tensor.rs | 22 +- cuda-oxide/cuteir/cute-rs/src/tile.rs | 6 +- cuda-oxide/cuteir/cute-rs/src/tiled_copy.rs | 10 + cuda-oxide/cuteir/cute-rs/src/tma.rs | 567 -- cuda-oxide/cuteir/cute-rs/src/tma/host.rs | 186 + cuda-oxide/cuteir/cute-rs/src/tma/mod.rs | 288 ++ cuda-oxide/cuteir/cute-rs/src/tma/tests.rs | 114 + cuda-oxide/cuteir/cute-rs/tests/arch.rs | 2 + .../compile_fail/global_store_storage.rs | 2 + .../compile_fail/global_store_storage.stderr | 6 +- .../tests/compile_fail/global_store_unsafe.rs | 2 + .../compile_fail/global_store_unsafe.stderr | 4 +- .../tests/compile_fail/immutable_store.rs | 3 + .../tests/compile_fail/immutable_store.stderr | 4 +- .../tests/compile_fail/invalid_storage.rs | 3 + .../tests/compile_fail/invalid_storage.stderr | 8 +- .../compile_fail/mma_operand_role_mismatch.rs | 3 + .../mma_operand_role_mismatch.stderr | 6 +- .../compile_fail/mma_scale_role_mismatch.rs | 3 + .../mma_scale_role_mismatch.stderr | 28 +- .../sm100_attention_invalid_layout.rs | 4 + .../sm100_attention_invalid_layout.stderr | 9 +- .../sm100_attention_operand_role.rs | 3 + .../sm100_attention_operand_role.stderr | 6 +- .../sm100_attention_tile_array_extent.rs | 5 + .../sm100_attention_tile_array_extent.stderr | 6 +- .../sm100_attention_tile_array_role.rs | 3 + .../sm100_attention_tile_array_role.stderr | 6 +- .../sm100_attention_tma_layout.rs | 5 + .../sm100_attention_tma_layout.stderr | 26 +- .../compile_fail/sm100_attention_ts_k16.rs | 4 + .../sm100_attention_ts_k16.stderr | 5 +- .../compile_fail/sm100_mma_shape_mismatch.rs | 6 + .../sm100_mma_shape_mismatch.stderr | 4 +- .../sm100_operand_role_mismatch.rs | 5 + .../sm100_operand_role_mismatch.stderr | 6 +- .../compile_fail/sm100_tma_layout_mismatch.rs | 4 + .../sm100_tma_layout_mismatch.stderr | 6 +- .../tests/compile_fail/tensor_mut_clone.rs | 5 + .../compile_fail/tensor_mut_clone.stderr | 8 +- .../tests/compile_fail/tensor_mut_copy.rs | 5 + .../tests/compile_fail/tensor_mut_copy.stderr | 8 +- .../tiled_copy_encoding_mismatch.rs | 3 + .../tiled_copy_encoding_mismatch.stderr | 4 +- .../compile_fail/tiled_copy_role_mismatch.rs | 3 + .../tiled_copy_role_mismatch.stderr | 4 +- ...tma_descriptor_prefetch_requires_unsafe.rs | 5 + ...descriptor_prefetch_requires_unsafe.stderr | 12 +- .../compile_fail/tma_store_layout_mismatch.rs | 6 + .../tma_store_layout_mismatch.stderr | 6 +- .../two_cta_ring_cursor_wrong_role.rs | 4 + .../two_cta_ring_cursor_wrong_role.stderr | 16 +- .../two_cta_ring_cursor_wrong_stages.rs | 4 + .../two_cta_ring_cursor_wrong_stages.stderr | 8 +- .../two_cta_ring_cursor_zero_stages.rs | 4 + .../two_cta_ring_cursor_zero_stages.stderr | 1 + .../two_cta_ring_non_tc_wrong_signal.rs | 5 + .../two_cta_ring_non_tc_wrong_signal.stderr | 24 +- .../two_cta_ring_policy_sealed.rs | 5 + .../two_cta_ring_policy_sealed.stderr | 8 +- .../compile_fail/two_cta_ring_wrong_signal.rs | 5 + .../two_cta_ring_wrong_signal.stderr | 20 +- .../compile_fail/two_cta_ring_zero_stages.rs | 5 + .../two_cta_ring_zero_stages.stderr | 1 + .../tests/compile_fail/warp_uniform_unsafe.rs | 2 + .../compile_fail/warp_uniform_unsafe.stderr | 4 +- .../tests/compile_fail/warp_uniform_width.rs | 2 + .../compile_fail/warp_uniform_width.stderr | 8 +- .../compile_pass/tma_descriptor_prefetch.rs | 5 + .../two_cta_ring_non_tc_receivers.rs | 5 + .../compile_pass/two_cta_ring_valid_cursor.rs | 4 + .../compile_pass/two_cta_ring_valid_signal.rs | 5 + .../cuteir/cute-rs/tests/global_store.rs | 4 + .../cuteir/cute-rs/tests/sm100_attention.rs | 15 +- .../cute-rs/tests/sm100_attention_pipeline.rs | 6 + .../cuteir/cute-rs/tests/type_safety.rs | 2 + .../cuteir/dialect-cute/src/attributes.rs | 1092 ---- .../dialect-cute/src/attributes/copy.rs | 45 + .../dialect-cute/src/attributes/epilogue.rs | 287 ++ .../dialect-cute/src/attributes/gemv.rs | 34 + .../dialect-cute/src/attributes/layout.rs | 122 + .../cuteir/dialect-cute/src/attributes/mma.rs | 353 ++ .../cuteir/dialect-cute/src/attributes/mod.rs | 49 + .../dialect-cute/src/attributes/pipeline.rs | 76 + .../dialect-cute/src/attributes/scheduler.rs | 69 + .../dialect-cute/src/attributes/tensor.rs | 187 + .../dialect-cute/src/attributes/tests.rs | 68 + .../cuteir/dialect-cute/src/copy_ops.rs | 4 + .../cuteir/dialect-cute/src/epilogue_ops.rs | 76 + .../cuteir/dialect-cute/src/gemm_tma_ops.rs | 37 + .../cuteir/dialect-cute/src/gemv_ops.rs | 40 + cuda-oxide/cuteir/dialect-cute/src/lib.rs | 2 + cuda-oxide/cuteir/dialect-cute/src/ops.rs | 1609 ------ .../cuteir/dialect-cute/src/ops/copy.rs | 382 ++ .../dialect-cute/src/ops/divisibility.rs | 114 + .../cuteir/dialect-cute/src/ops/ldmatrix.rs | 170 + cuda-oxide/cuteir/dialect-cute/src/ops/mod.rs | 28 + .../cuteir/dialect-cute/src/ops/tests.rs | 789 +++ cuda-oxide/cuteir/dialect-cute/src/ops/tma.rs | 277 + .../cuteir/dialect-cute/src/pipeline_ops.rs | 67 + .../cuteir/dialect-cute/src/scheduler_ops.rs | 31 + .../dialect-cute/src/sm100_attention_ops.rs | 97 +- .../cuteir/dialect-cute/src/sm100_ops.rs | 207 +- .../cuteir/dialect-cute/src/smem_mma_ops.rs | 86 + .../cuteir/dialect-cute/src/tensor_ops.rs | 27 + cuda-oxide/cuteir/dialect-cute/src/types.rs | 1220 ----- .../cuteir/dialect-cute/src/types/epilogue.rs | 103 + .../cuteir/dialect-cute/src/types/gemv.rs | 203 + .../cuteir/dialect-cute/src/types/mma.rs | 217 + .../cuteir/dialect-cute/src/types/mod.rs | 56 + .../cuteir/dialect-cute/src/types/pipeline.rs | 105 + .../dialect-cute/src/types/scheduler.rs | 51 + .../cuteir/dialect-cute/src/types/tensor.rs | 317 ++ .../cuteir/dialect-cute/src/types/tests.rs | 232 + .../cuteir/dialect-cute/src/types/tma.rs | 140 + cuda-oxide/cuteir/dialect-cute/src/verify.rs | 4548 ----------------- .../cuteir/dialect-cute/src/verify/cfg.rs | 286 ++ .../dialect-cute/src/verify/epilogue.rs | 691 +++ .../cuteir/dialect-cute/src/verify/gemv.rs | 212 + .../cuteir/dialect-cute/src/verify/ghosts.rs | 304 ++ .../cuteir/dialect-cute/src/verify/mma.rs | 782 +++ .../cuteir/dialect-cute/src/verify/mod.rs | 312 ++ .../dialect-cute/src/verify/pipeline.rs | 607 +++ .../dialect-cute/src/verify/provenance.rs | 533 ++ .../dialect-cute/src/verify/scheduler.rs | 382 ++ .../cuteir/dialect-cute/src/verify/tensor.rs | 131 + .../cuteir/dialect-cute/src/verify/tests.rs | 754 +++ .../cuteir/dialect-cute/src/verify/tma.rs | 137 + .../cuteir/dialect-cute/tests/global_store.rs | 2 + .../dialect-cute/tests/sm100_attention.rs | 10 + cuda-oxide/cuteir/docs/translation.md | 6 +- .../blockscale_gemm_ctm/src/kernel.rs | 39 +- .../examples/blockscale_gemm_ctm/src/main.rs | 8 + .../blockscale_gemm_ctm/src/physical.rs | 10 + .../examples/blockscale_gemm_cute/src/main.rs | 43 + .../examples/elementwise_cute/README.md | 2 +- .../examples/elementwise_cute/src/main.rs | 8 + .../fmha_decode_sm120a_cute/Cargo.lock | 1 + .../fmha_decode_sm120a_cute/Cargo.toml | 2 + .../fmha_decode_sm120a_cute/README.md | 7 + .../reference/ctm-migration-results.json | 283 + .../fmha_decode_sm120a_cute/src/config.rs | 49 + .../fmha_decode_sm120a_cute/src/kernel.rs | 44 +- .../fmha_decode_sm120a_cute/src/main.rs | 5 + .../fmha_decode_sm120a_cute/src/runner.rs | 3 + .../fmha_prefill_sm120a_cute/Cargo.lock | 1 + .../fmha_prefill_sm120a_cute/Cargo.toml | 2 + .../fmha_prefill_sm120a_cute/README.md | 7 + .../reference/ctm-migration-results.json | 258 + .../src/comparison.rs | 15 +- .../fmha_prefill_sm120a_cute/src/fixture.rs | 27 + .../fmha_prefill_sm120a_cute/src/kernel.rs | 74 +- .../fmha_prefill_sm120a_cute/src/main.rs | 5 + .../fmha_prefill_sm120a_cute/src/runner.rs | 35 + .../fp16_gemm_256x352_ctm/src/kernel.rs | 37 +- .../fp16_gemm_256x352_ctm/src/main.rs | 34 + .../fp16_gemm_256x352_ctm/src/physical.rs | 5 + .../examples/fp16_gemm_256x352_cute/README.md | 2 +- .../fp16_gemm_256x352_cute/src/kernel.rs | 20 + .../fp16_gemm_256x352_cute/src/main.rs | 37 + .../examples/nvfp4_gemv_cute/src/main.rs | 40 + .../sm100_fmha_prefill_cute/src/kernel.rs | 56 + .../sm100_fmha_prefill_cute/src/main.rs | 24 + .../reference/attention-feature-gaps.md | 10 +- .../reference/measurement-notes.md | 2 +- .../src/comparison.rs | 14 + .../sm100_mla_decode_fp16_cute/src/kernel.rs | 51 + .../sm100_mla_decode_fp16_cute/src/main.rs | 36 + .../scripts/check-primitives-api-compile.py | 116 +- 385 files changed, 27547 insertions(+), 20695 deletions(-) create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/common.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/common_sm100.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/cooperative.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/emit.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/mod.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/registry.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/sm100_attention.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/tma.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/emit.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/mod.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/sm100.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/sm120.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/independent_sm100.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/ldmatrix.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/mod.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm100.rs rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{cute_sm100_attention.rs => cute/mma/sm100_attention.rs} (60%) create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm120.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mod.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/pipeline/mod.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/pipeline/sm100.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/pipeline/sm120.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/registry.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/scheduler/mod.rs create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/scheduler/sm120.rs rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{cute_sm100.rs => cute/sm100.rs} (52%) rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{cute.rs => cute/tensor/elementwise.rs} (92%) rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{cute_gemv.rs => cute/tensor/gemv.rs} (91%) create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/mod.rs rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{cute_attention.rs => cute/tests/cooperative.rs} (59%) create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tests/gemm.rs delete mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_gemm.rs rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{mir_checked.rs => mir/checked.rs} (91%) rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{mir_core.rs => mir/core.rs} (92%) rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{mir_memory.rs => mir/memory.rs} (93%) create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/mod.rs rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{mir_shared.rs => mir/shared.rs} (95%) rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{nvvm_attention.rs => nvvm/attention.rs} (93%) rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{nvvm_cluster.rs => nvvm/cluster.rs} (87%) create mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/mod.rs rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{nvvm_numeric.rs => nvvm/numeric.rs} (88%) rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{nvvm_sm120.rs => nvvm/sm120.rs} (95%) rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{nvvm_sregs.rs => nvvm/sregs.rs} (94%) rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{nvvm_sync_restrict.rs => nvvm/sync_restrict.rs} (89%) rename cuda-oxide/crates/cuda-oxide-mlir-export/src/{nvvm_tcgen05.rs => nvvm/tcgen05.rs} (91%) delete mode 100644 cuda-oxide/crates/cuda-oxide-mlir-export/stack-temporaries.md rename cuda-oxide/crates/mir-importer/src/translator/cute/{block_scaled_emit.rs => block_scaled/emit.rs} (93%) rename cuda-oxide/crates/mir-importer/src/translator/cute/{block_scaled.rs => block_scaled/mod.rs} (79%) create mode 100644 cuda-oxide/crates/mir-importer/src/translator/cute/classify.rs create mode 100644 cuda-oxide/crates/mir-importer/src/translator/cute/dispatch.rs rename cuda-oxide/crates/mir-importer/src/translator/cute/{epilogue_emit.rs => epilogue/emit.rs} (92%) rename cuda-oxide/crates/mir-importer/src/translator/cute/{epilogue.rs => epilogue/mod.rs} (79%) create mode 100644 cuda-oxide/crates/mir-importer/src/translator/cute/identity.rs rename cuda-oxide/crates/mir-importer/src/translator/cute/{pipeline_emit.rs => pipeline/emit.rs} (91%) rename cuda-oxide/crates/mir-importer/src/translator/cute/{pipeline.rs => pipeline/mod.rs} (85%) create mode 100644 cuda-oxide/crates/mir-importer/src/translator/cute/registry.rs rename cuda-oxide/crates/mir-importer/src/translator/cute/{scheduler_emit.rs => scheduler/emit.rs} (94%) rename cuda-oxide/crates/mir-importer/src/translator/cute/{scheduler.rs => scheduler/mod.rs} (83%) rename cuda-oxide/crates/mir-importer/src/translator/cute/{sm100_emit.rs => sm100.rs} (85%) rename cuda-oxide/crates/mir-importer/src/translator/cute/{sm100_attention_emit.rs => sm100_attention.rs} (89%) rename cuda-oxide/crates/mir-importer/src/translator/cute/{smem_mma_emit.rs => smem_mma/emit.rs} (93%) rename cuda-oxide/crates/mir-importer/src/translator/cute/{smem_mma.rs => smem_mma/mod.rs} (62%) rename cuda-oxide/crates/mir-importer/src/translator/cute/{tensor_emit.rs => tensor/emit.rs} (95%) rename cuda-oxide/crates/mir-importer/src/translator/cute/{tensor.rs => tensor/mod.rs} (84%) create mode 100644 cuda-oxide/cuteir/ctm-rs/src/async_copy.rs create mode 100644 cuda-oxide/cuteir/ctm-rs/src/convert.rs create mode 100644 cuda-oxide/cuteir/ctm-rs/src/float.rs create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/attention_api.rs create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_async_wait_runtime.rs create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_async_wait_runtime.stderr create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_f16_mma_roles.rs create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_f16_mma_roles.stderr create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_pass/attention_primitives.rs delete mode 100644 cuda-oxide/cuteir/cute-layout/src/validation.rs create mode 100644 cuda-oxide/cuteir/cute-layout/src/validation/cooperative.rs create mode 100644 cuda-oxide/cuteir/cute-layout/src/validation/error.rs create mode 100644 cuda-oxide/cuteir/cute-layout/src/validation/ldmatrix.rs create mode 100644 cuda-oxide/cuteir/cute-layout/src/validation/mod.rs create mode 100644 cuda-oxide/cuteir/cute-layout/src/validation/tests.rs create mode 100644 cuda-oxide/cuteir/cute-layout/src/validation/thread_value.rs create mode 100644 cuda-oxide/cuteir/cute-layout/src/validation/tma.rs delete mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled/__compiler.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled/gemv.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled/layouts.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled/mod.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled/scales.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled/tensor.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled/tests.rs delete mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled_mma.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/__compiler.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/epilogue.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/fragments.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/mod.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/tests.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/tiled.rs delete mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100/__compiler.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100/mma.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100/mod.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100/pipeline.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100/shared.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100/tests.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100/tma.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100/tmem.rs delete mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100_attention.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100_attention/__compiler.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100_attention/mma.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100_attention/mod.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100_attention/shared.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100_attention/tma.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/sm100_attention/tmem.rs delete mode 100644 cuda-oxide/cuteir/cute-rs/src/tma.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/tma/host.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/tma/mod.rs create mode 100644 cuda-oxide/cuteir/cute-rs/src/tma/tests.rs delete mode 100644 cuda-oxide/cuteir/dialect-cute/src/attributes.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/attributes/copy.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/attributes/epilogue.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/attributes/gemv.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/attributes/layout.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/attributes/mma.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/attributes/mod.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/attributes/pipeline.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/attributes/scheduler.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/attributes/tensor.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/attributes/tests.rs delete mode 100644 cuda-oxide/cuteir/dialect-cute/src/ops.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/ops/copy.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/ops/divisibility.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/ops/ldmatrix.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/ops/mod.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/ops/tests.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/ops/tma.rs delete mode 100644 cuda-oxide/cuteir/dialect-cute/src/types.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/types/epilogue.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/types/gemv.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/types/mma.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/types/mod.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/types/pipeline.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/types/scheduler.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/types/tensor.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/types/tests.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/types/tma.rs delete mode 100644 cuda-oxide/cuteir/dialect-cute/src/verify.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/verify/cfg.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/verify/epilogue.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/verify/gemv.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/verify/ghosts.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/verify/mma.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/verify/mod.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/verify/pipeline.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/verify/provenance.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/verify/scheduler.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/verify/tensor.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/verify/tests.rs create mode 100644 cuda-oxide/cuteir/dialect-cute/src/verify/tma.rs create mode 100644 cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/reference/ctm-migration-results.json create mode 100644 cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/reference/ctm-migration-results.json diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/README.md b/cuda-oxide/crates/cuda-oxide-mlir-export/README.md index 4c94df5643..c4ced98c17 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/README.md +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/README.md @@ -1,6 +1,6 @@ -# CUDA Oxide MLIR export +# cuda-oxide MLIR export -This crate maps CUDA Oxide operations to MLIR accepted by the pinned CUTLASS +This crate maps cuda-oxide operations to MLIR accepted by the pinned CUTLASS compiler. It uses `pliron-mlir-export` to build and render the MLIR text: ```text @@ -13,8 +13,7 @@ registry + diagnostics pinned consumer profiles The generic crate can translate any Pliron dialect. CUDA-specific mappings and compiler-version requirements live here. -The profile targets the official CUTLASS 4.7 compiler library. CUDA -Oxide installs the exact NVIDIA release archive, verifies both the archive and +The profile targets the official CUTLASS 4.7 compiler library. cuda-oxide installs the exact NVIDIA release archive, verifies both the archive and library digests, and loads that versioned library explicitly: ```text @@ -35,7 +34,7 @@ shared high-level Pliron module The mapping packs translate CuTe operations after shared MIR preparation. The pinned library compiles the resulting MLIR, and the backend extracts and -validates its cubin. CUDA Oxide embeds and launches that cubin through its +validates its cubin. cuda-oxide embeds and launches that cubin through its ordinary artifact loader. Set `CUDA_OXIDE_MLIR_OUTPUT=` to keep the MLIR text for inspection; select `CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir` to compile it. @@ -52,7 +51,8 @@ Unsupported primitive selectors and intrinsic identities fail at export. 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` only. The default full +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, @@ -61,8 +61,8 @@ semantics. Explicit NVVM shared offsets preserve AS3/AS7; Rust address casts convert named pointers to generic addresses first. See the [coverage and validation contract](../../cuteir/ctm-rs/docs/primitives-coverage.md). -See [why stack temporaries are moved to the function entry](stack-temporaries.md) -for a short explanation of enum scratch allocation placement. +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 @@ -194,3 +194,18 @@ 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..e00dac0375 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/common.rs @@ -0,0 +1,427 @@ +/* + * 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 dialect_mir::ops::{MirExternSharedOp, MirInsertFieldOp}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, + r#type::TypeHandle, +}; +use pliron_mlir_export::{ + MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, OperationInput, + OperationTranslation, TranslationRegistry, TranslationSession, TypeTranslation, +}; + +/// 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; + 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..d06eaa74b0 --- /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 pliron_mlir_export::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..3d4814088a --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/cooperative.rs @@ -0,0 +1,135 @@ +/* + * 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 cute_layout::{IntTuple, validate_cooperative_copy_plan}; +use dialect_cute::ops::CuteCopyG2SOp; +use dialect_mir::types::MirFP16Type; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, + r#type::TypeHandle, +}; +use pliron_mlir_export::{ + MlirAttribute, MlirOperation, MlirType, OperationInput, OperationTranslation, + TranslationSession, +}; + +/// Emit cooperative FP16 global-to-shared 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(); + if !elem.deref(ctx).is::() { + return Err("CuTe attention copy supports f16 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, + MlirType::Float(pliron_mlir_export::MlirFloatType::F16), + )?; + 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..482c6fa010 --- /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 cute_layout::{ComposedLayout, Layout, Swizzle}; +use pliron_mlir_export::{ + MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, OperationInput, + TranslationSession, +}; + +/// 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.leaves(); + let strides = layout.stride.leaves(); + 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 inner row/column layout, fixed offset, and outer 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.inner(), row, column)?; + let offset = self.constant(layout.offset(), width)?; + let plain = self.binary("arith.addi", plain, offset)?; + self.swizzle(layout.outer(), 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.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.shift.unsigned_abs()), width)?; + let moved = self.binary( + if swizzle.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..f19d29d1f9 --- /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 dialect_cute::{ + attributes::{ + CuteComposedLayoutAttr, CuteCopyAtomAttr, CuteDivisibilityAttr, CuteLayoutAttr, + CuteMatrixRoleAttr, + }, + ops::{CuteCopyG2SOp, CuteLdmatrixOp}, +}; +use pliron_mlir_export::{DropAttribute, TranslationError, TranslationRegistry}; + +/// 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..1e122fd59b --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/sm100_attention.rs @@ -0,0 +1,91 @@ +/* + * 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 dialect_cute::sm100_attention_ops::*; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; +use pliron_mlir_export::{ + MlirOperation, MlirType, OperationInput, OperationTranslation, TranslationSession, +}; + +/// 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..4de8655b04 --- /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 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}, +}; +use pliron_mlir_export::{ + MlirAttribute, MlirLocation, MlirOperation, MlirType, MlirValueUse, OperationInput, + OperationTranslation, TranslationRegistry, TranslationSession, TypeTranslation, +}; + +/// 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.inner().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..ea4a0f664e --- /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 dialect_cute::attributes::CuteComposedLayoutAttr; +use pliron_mlir_export::{ + MlirAttribute, MlirBlock, MlirLocation, MlirOperation, MlirRegion, MlirResult, MlirType, + MlirValueUse, TranslationSession, +}; + +/// 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.outer(); + let ptr = self.cute_ptr( + ptr, + "f16", + "smem", + 128, + Some(format!( + "S<{},{},{}>", + swizzle.bits, + swizzle.base + 1, + swizzle.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..364453748b --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/sm100.rs @@ -0,0 +1,88 @@ +/* + * 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 dialect_cute::sm100_ops::*; +use pliron_mlir_export::{MlirAttribute, MlirFloatType, MlirType, MlirValueUse}; + +/// 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.outer(); + let mask = e.constant(32, (((1u32 << swizzle.bits) - 1) << swizzle.base) as i128)?; + let shift = e.constant(32, swizzle.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..f92cb606e2 --- /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 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, +}; +use pliron_mlir_export::{ + MlirAttribute, MlirLocation, MlirOperation, MlirType, MlirValueUse, OperationInput, + OperationTranslation, TranslationSession, +}; + +/// 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(pliron_mlir_export::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(pliron_mlir_export::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..1f181caf17 --- /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 dialect_cute::sm100_attention_ops::*; +use pliron_mlir_export::{DropAttribute, TranslationError, TranslationRegistry}; + +/// 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..4f44ad1a30 --- /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 cute_layout::{ComposedLayout, IntTuple, OffsetUnit}; +use dialect_cute::{attributes::CuteMatrixRoleAttr, ops::CuteLdmatrixOp}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; +use pliron_mlir_export::{ + MlirAttribute, MlirOperation, MlirType, OperationInput, OperationTranslation, + TranslationSession, +}; + +/// 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.inner().modes().len() != 2 + || !layout + .inner() + .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 + .inner() + .shape + .leaves() + .iter() + .zip(&layout.inner().stride.leaves()) + { + 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.outer(); + 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.inner().checked_size()? > cute_layout::validation::MAX_STATIC_VALIDATION_ELEMENTS + { + return None; + } + let modes = layout.inner().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.outer(); + 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.inner(), 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.inner(), 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.outer(), low)?; + let local = e.swizzle(layout.outer(), 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..41101f4d2b --- /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::cute::emit::{Emit, llvm_ptr_type, property, segments}; +use dialect_cute::sm100_ops::*; +use pliron_mlir_export::MlirValueUse; + +/// 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_sm100_attention.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm100_attention.rs similarity index 60% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_sm100_attention.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm100_attention.rs index cd1576cb47..c6606877c9 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_sm100_attention.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm100_attention.rs @@ -2,8 +2,13 @@ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ -//! Independent layout-aware SM100 MMA and tiled TMA operations. -use crate::cute_sm100::{Emit, integer_property, llvm_ptr_type, property, segments}; + +//! Independent SM100 attention MMA and layout-derived shared descriptors. + +use crate::cute::{ + common_sm100::contract, + emit::{Emit, llvm_ptr_type, property, segments}, +}; use dialect_cute::{attributes::CuteComposedLayoutAttr, sm100_attention_ops::*}; use pliron::{ common_traits::Verify, @@ -11,34 +16,14 @@ use pliron::{ operation::Operation, }; use pliron_mlir_export::{ - DropAttribute, MlirOperation, MlirType, MlirValueUse, OperationInput, OperationTranslation, - TranslationError, TranslationRegistry, TranslationSession, + MlirOperation, MlirType, MlirValueUse, OperationInput, OperationTranslation, TranslationSession, }; -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(()) -} -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(()) -} -struct Mma(bool); +/// 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, @@ -148,6 +133,7 @@ fn descriptor_offset(rows: u32, k_extent: u32, mn: bool, k: u32) -> u32 { ((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, @@ -202,69 +188,3 @@ fn descriptor( let offset = e.binary("arith.shrui", offset, shift)?; e.binary("arith.addi", base, offset) } -struct Tma(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/mma/sm120.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm120.rs new file mode 100644 index 0000000000..a8fdc1f52d --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm120.rs @@ -0,0 +1,1083 @@ +/* + * 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 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}, +}; +use pliron_mlir_export::{ + MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, OperationInput, + OperationTranslation, TranslationSession, +}; + +/// 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> { + CuteTiledGemmOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + 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(), + &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 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, + 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 = cast_int( + session, + operations, + scale_a, + MlirType::Integer(16), + "arith.trunci", + location, + )?; + let scale_b = 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(pliron_mlir_export::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(pliron_mlir_export::MlirFloatType::Other( + "f4E2M1FN".into(), + ))), + ); + operation_.properties.insert( + "b_type".into(), + MlirAttribute::Type(MlirType::Float(pliron_mlir_export::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(pliron_mlir_export::MlirFloatType::Other( + "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: 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..2196c21c65 --- /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 dialect_cute::sm100_ops::*; +use pliron_mlir_export::{MlirAttribute, MlirRegion, MlirType, MlirValueUse}; + +/// 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..2b0d1d0e7b --- /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 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}, +}; +use pliron_mlir_export::{ + MlirAttribute, MlirBlock, MlirLocation, MlirOperation, MlirRegion, MlirResult, MlirType, + MlirValueUse, OperationInput, OperationTranslation, TranslationRegistry, TranslationSession, + TypeTranslation, +}; + +/// 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..afde1a03a4 --- /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 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}; +use pliron_mlir_export::{DropAttribute, TranslationError, TranslationRegistry}; + +/// 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..059daf42cb --- /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 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}, +}; +use pliron_mlir_export::{ + MlirAttribute, MlirOperation, MlirType, OperationInput, OperationTranslation, + TranslationRegistry, TranslationSession, TypeTranslation, +}; + +/// 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 similarity index 52% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_sm100.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/sm100.rs index d90ba2eab1..05257428bd 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_sm100.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/sm100.rs @@ -16,21 +16,28 @@ //! second partition -> last use of A collector //! ``` -use dialect_cute::{ - attributes::{CuteComposedLayoutAttr, CuteTmaStorePipelineAttr}, - sm100_ops::*, +#[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, }; +#[cfg(test)] +use dialect_cute::attributes::CuteComposedLayoutAttr; +use dialect_cute::{attributes::CuteTmaStorePipelineAttr, sm100_ops::*}; use pliron::{ common_traits::Verify, context::{Context, Ptr}, operation::Operation, }; use pliron_mlir_export::{ - DropAttribute, MlirAttribute, MlirBlock, MlirFloatType, MlirLocation, MlirOperation, - MlirRegion, MlirResult, MlirType, MlirValueUse, OperationInput, OperationTranslation, + DropAttribute, MlirAttribute, MlirOperation, MlirType, OperationInput, OperationTranslation, TranslationError, TranslationRegistry, TranslationSession, }; +/// 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> { @@ -39,6 +46,7 @@ pub(crate) fn register_cute_sm100_pack( 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 { @@ -75,38 +83,77 @@ pub(crate) fn register_cute_sm100_pack( Ok(()) } +/// Select the runtime action after the source plan has been verified. #[derive(Clone, Copy)] -enum Kind { +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 { - Tmem(CuteSm100TmemPlanAttr), - Mma(CuteSm100MmaPlanAttr), - Epilogue(CuteSm100EpiloguePlanAttr), - Copy(CuteSm100ClusterCopyPlanAttr), - Pipeline(CuteSm100PipelinePlanAttr), - Store(CuteTmaStorePipelineAttr), + /// 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, } @@ -233,439 +280,10 @@ impl OperationTranslation for Recipe { } } -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(()) -} - -fn instruction_descriptor(m: u32, n: u32) -> u32 { - (1 << 4) | ((n >> 3) << 17) | ((m >> 4) << 24) -} - -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.outer(); - let mask = e.constant(32, (((1u32 << swizzle.bits) - 1) << swizzle.base) as i128)?; - let shift = e.constant(32, swizzle.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(()) -} - -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(()) -} - -pub(crate) struct Emit<'a, 'b> { - pub(crate) session: &'a mut TranslationSession<'b>, - pub(crate) location: MlirLocation, - pub(crate) ops: Vec, -} -impl Emit<'_, '_> { - 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) - } - pub(crate) fn fresh(&mut self, ty: MlirType) -> (MlirResult, MlirValueUse) { - let id = self.session.fresh_value(); - (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) - } - fn emit(&mut self, name: &str, operands: Vec) -> Result<(), String> { - let op = self.operation(name, operands)?; - self.ops.push(op); - Ok(()) - } - 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) - } - pub(crate) fn cast( - &mut self, - name: &str, - operand: MlirValueUse, - ty: MlirType, - ) -> Result { - self.value(name, vec![operand], ty) - } - 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) - } - 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) - } - 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) - } - } - 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}>" - ))?, - ) - } - 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) - } - 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), - } - } - 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.outer(); - let ptr = self.cute_ptr( - ptr, - "f16", - "smem", - 128, - Some(format!( - "S<{},{},{}>", - swizzle.bits, - swizzle.base + 1, - swizzle.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), - ) - } - 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(()) - } - fn region(&mut self, operations: Vec) -> MlirRegion { - MlirRegion { - blocks: vec![MlirBlock { - id: self.session.fresh_block(), - arguments: vec![], - operations, - }], - } - } - 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(()) - } -} -fn dialect(s: &str) -> Result { - MlirType::dialect(s) -} -pub(crate) fn llvm_ptr_type(space: u32) -> Result { - dialect(&if space == 0 { - "!llvm.ptr".into() - } else { - format!("!llvm.ptr<{space}>") - }) -} -pub(crate) fn property(op: &mut MlirOperation, key: &str, value: &str) -> Result<(), String> { - op.properties - .insert(key.into(), MlirAttribute::dialect(value)?); - Ok(()) -} -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), - }, - ); -} -pub(crate) fn segments(op: &mut MlirOperation, values: &[i32]) { - op.properties.insert( - "operandSegmentSizes".into(), - MlirAttribute::DenseI32Array(values.to_vec()), - ); -} - #[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}; @@ -686,6 +304,7 @@ mod tests { 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( @@ -697,6 +316,7 @@ mod tests { .unwrap(), ) } + /// Build the paired-MMA fixture plan for two independently selected N partitions. fn mma_plan(n0: u32, n1: u32) -> CuteSm100MmaPlanAttr { CuteSm100MmaPlanAttr { m: 256, @@ -711,6 +331,7 @@ mod tests { 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); @@ -747,6 +368,7 @@ mod tests { 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, @@ -770,6 +392,7 @@ mod tests { .unwrap(); pliron_mlir_export::render_module(&target) } + /// 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)] { @@ -810,6 +433,7 @@ mod tests { 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(); @@ -840,6 +464,7 @@ mod tests { 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(); @@ -876,6 +501,7 @@ mod tests { ); } + /// 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(); diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/elementwise.rs similarity index 92% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/cute.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/elementwise.rs index 45f852faf5..6eb5295545 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/elementwise.rs @@ -82,6 +82,7 @@ pub(crate) fn register_cute_elementwise_mappings( Ok(()) } +/// Preserve reviewed elementwise view layouts as CuTe memref types. pub(crate) struct TensorViewTypeTranslation; impl TypeTranslation for TensorViewTypeTranslation { @@ -107,6 +108,7 @@ impl TypeTranslation for TensorViewTypeTranslation { } } +/// Build a contiguous global CuTe view with an explicit LLVM-to-CuTe pointer bridge. struct TensorMakeTranslation; impl OperationTranslation for TensorMakeTranslation { @@ -177,6 +179,7 @@ impl OperationTranslation for TensorMakeTranslation { } } +/// Divide the verified one-dimensional view into statically sized CuTe tiles. struct TensorZippedDivideTranslation; impl OperationTranslation for TensorZippedDivideTranslation { @@ -208,6 +211,7 @@ impl OperationTranslation for TensorZippedDivideTranslation { } } +/// Select one tile using its runtime coordinate and verified layout. struct TensorSliceTranslation; impl OperationTranslation for TensorSliceTranslation { @@ -244,6 +248,7 @@ impl OperationTranslation for TensorSliceTranslation { } } +/// Recover an absolute element index from the verified view's producer chain. struct TensorBaseTranslation; impl OperationTranslation for TensorBaseTranslation { @@ -264,6 +269,7 @@ impl OperationTranslation for TensorBaseTranslation { } } +/// Check that the selected tile lies wholly within the original logical length. struct TensorIsFullTranslation; impl OperationTranslation for TensorIsFullTranslation { @@ -331,6 +337,7 @@ impl OperationTranslation for TensorIsFullTranslation { } } +/// Copy a full global tile into its register fragment. struct TensorLoadIntoTranslation; impl OperationTranslation for TensorLoadIntoTranslation { @@ -361,6 +368,7 @@ impl OperationTranslation for TensorLoadIntoTranslation { } } +/// Copy a register fragment into a full global tile. struct TensorStoreFromTranslation; impl OperationTranslation for TensorStoreFromTranslation { @@ -391,6 +399,7 @@ impl OperationTranslation for TensorStoreFromTranslation { } } +/// Store one element only when its absolute index is within the original view. struct TensorStoreElementAbsTranslation; impl OperationTranslation for TensorStoreElementAbsTranslation { @@ -450,14 +459,20 @@ impl OperationTranslation for TensorStoreElementAbsTranslation { } } +/// 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, @@ -500,6 +515,7 @@ fn resolve_tensor_view( )) } +/// 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 { @@ -512,6 +528,7 @@ fn tensor_tile_state(ctx: &Context, value: Value) -> Result, state: &TensorViewState, @@ -591,12 +608,16 @@ fn emit_saturating_tile_base( 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, @@ -741,6 +762,7 @@ fn emit_full_tile_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!( @@ -774,6 +796,7 @@ fn require_elementwise_view(view: &CuteTensorViewType) -> Result<(), String> { 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); @@ -783,6 +806,7 @@ fn source_view(ctx: &Context, value: Value) -> Result Result<&'static str, String> { match space { CuteTensorAddressSpaceAttr::Gmem => Ok("gmem"), @@ -792,6 +816,7 @@ fn address_space_spelling(space: CuteTensorAddressSpaceAttr) -> Result<&'static } } +/// Spell a reviewed elementwise layout in CUTLASS syntax. fn layout_spelling(layout: CuteTensorLayoutAttr) -> Result { match layout { CuteTensorLayoutAttr::Contiguous1D => Ok("?{i64}:1".into()), @@ -805,6 +830,7 @@ fn layout_spelling(layout: CuteTensorLayoutAttr) -> Result { } } +/// 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}")), @@ -816,6 +842,7 @@ fn scalar_spelling(ty: &MlirType) -> Result { } } +/// 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}>>")) } @@ -860,6 +887,7 @@ fn bridge_llvm_pointer_to_cute( 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!( @@ -882,6 +910,7 @@ fn llvm_pointer_address_space(ty: &MlirType) -> Result { }) } +/// 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") @@ -890,6 +919,7 @@ fn llvm_pointer_type(address_space: u32) -> Result { } } +/// Build the CuTe memref type for a register-copy layout and its storage element. fn cute_memref_type( element: &str, space: &str, @@ -901,6 +931,7 @@ fn cute_memref_type( )) } +/// Require exact result/operand counts and no leftover attributes, regions, or successors. fn expect_shape( name: &str, input: &OperationInput, @@ -922,6 +953,7 @@ fn expect_shape( Ok(()) } +/// Build a target operation with the supplied operands, results, and source location. fn operation( name: &str, results: Vec, @@ -935,11 +967,13 @@ fn operation( 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, @@ -956,6 +990,7 @@ fn integer_constant( Ok(operation) } +/// Build an arith.cmpi using the explicit target predicate number. fn integer_compare( result: MlirResult, left: MlirValueUse, @@ -974,6 +1009,7 @@ fn integer_compare( Ok(operation) } +/// Build a two-operand target operation using the supplied result. fn binary( name: &str, result: MlirResult, @@ -984,6 +1020,7 @@ fn binary( 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, @@ -993,14 +1030,17 @@ fn unary( 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( diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_gemv.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/gemv.rs similarity index 91% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_gemv.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/gemv.rs index 866c1d0822..4d6eb615b3 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_gemv.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/gemv.rs @@ -21,7 +21,9 @@ //! The mapping translates CuTe operations directly into operations accepted //! by CUTLASS. -use crate::cute::{TensorViewTypeTranslation, register_cute_elementwise_mappings}; +use crate::cute::tensor::elementwise::{ + TensorViewTypeTranslation, register_cute_elementwise_mappings, +}; use dialect_cute::{ attributes::{ @@ -46,12 +48,19 @@ use pliron_mlir_export::{ TranslationSession, TypeTranslation, }; +/// 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. @@ -67,6 +76,7 @@ pub(crate) fn register_cute_cutlass_profile_pack( 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)?; @@ -80,6 +90,7 @@ fn register_gemv_mappings(registry: &mut TranslationRegistry) -> Result<(), Tran Ok(()) } +/// Dispatch verified tensor views to their GEMV carrier or elementwise CuTe representation. struct CombinedTensorViewTypeTranslation; impl TypeTranslation for CombinedTensorViewTypeTranslation { @@ -101,6 +112,7 @@ impl TypeTranslation for CombinedTensorViewTypeTranslation { } } +/// Represent a verified scaled view with its runtime pointer and shape carrier. struct ScaledViewTypeTranslation; impl TypeTranslation for ScaledViewTypeTranslation { @@ -119,6 +131,7 @@ impl TypeTranslation for ScaledViewTypeTranslation { } } +/// Represent a verified GEMV fragment as the ordered FP32 lane carrier. struct FragmentTypeTranslation; impl TypeTranslation for FragmentTypeTranslation { @@ -137,6 +150,7 @@ impl TypeTranslation for FragmentTypeTranslation { } } +/// Pack runtime pointer, rows, columns, and stride into the two-dimensional view carrier. struct TensorMake2DTranslation; impl OperationTranslation for TensorMake2DTranslation { @@ -162,6 +176,7 @@ impl OperationTranslation for TensorMake2DTranslation { } } +/// Combine values and scales into the verified full scaled-view carrier. struct ScaledViewMakeTranslation; impl OperationTranslation for ScaledViewMakeTranslation { @@ -201,6 +216,7 @@ impl OperationTranslation for ScaledViewMakeTranslation { } } +/// Select the runtime row while retaining scale-layout and shape dependencies. struct ScaledViewRowTranslation; impl OperationTranslation for ScaledViewRowTranslation { @@ -239,6 +255,7 @@ impl OperationTranslation for ScaledViewRowTranslation { } } +/// Select the source row's 64-element K tile and corresponding scale coordinates. struct ScaledViewKTileTranslation; impl OperationTranslation for ScaledViewKTileTranslation { @@ -267,6 +284,7 @@ impl OperationTranslation for ScaledViewKTileTranslation { } } +/// Load one verified K64 tile and convert packed values/scales into FP32 fragments. struct ScaledViewLoadTranslation; impl OperationTranslation for ScaledViewLoadTranslation { @@ -309,6 +327,7 @@ impl OperationTranslation for ScaledViewLoadTranslation { } } +/// Accumulate the K64 fragment product in the source's explicit FP32 order. struct DotTranslation; impl OperationTranslation for DotTranslation { @@ -345,6 +364,7 @@ impl OperationTranslation for DotTranslation { } } +/// 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, @@ -649,12 +669,16 @@ fn emit_scaled_k64_load( 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, } +/// Build a narrow-format CuTe memref and load its verified contiguous register vector. fn emit_narrow_load( session: &mut TranslationSession<'_>, operations: &mut Vec, @@ -783,6 +807,7 @@ fn emit_narrow_load( Ok(extended) } +/// Emit the 64 FP32 products and additions in source order without reassociation. fn emit_ordered_k64_dot( session: &mut TranslationSession<'_>, matrix: MlirValueUse, @@ -891,6 +916,7 @@ fn emit_ordered_k64_dot( Ok(operations) } +/// Compute an unsigned ceiling quotient without overflowing an add-before-divide expression. fn emit_ceil_div( session: &mut TranslationSession<'_>, operations: &mut Vec, @@ -941,6 +967,7 @@ fn emit_ceil_div( ) } +/// Append an i64 constant while preserving the unsigned input's bit pattern. fn emit_i64_constant( session: &mut TranslationSession<'_>, operations: &mut Vec, @@ -952,6 +979,7 @@ fn emit_i64_constant( Ok(use_) } +/// Append a two-operand operation whose result has the left operand's type. fn emit_binary( session: &mut TranslationSession<'_>, operations: &mut Vec, @@ -971,6 +999,7 @@ fn emit_binary( Ok(use_) } +/// Extract one compile-time vector lane into a fresh scalar result. fn emit_vector_extract( session: &mut TranslationSession<'_>, operations: &mut Vec, @@ -1000,6 +1029,7 @@ fn emit_vector_extract( Ok(use_) } +/// Insert ordered runtime fields into the requested LLVM struct result. fn build_aggregate( session: &mut TranslationSession<'_>, final_result: MlirResult, @@ -1035,6 +1065,7 @@ fn build_aggregate( Ok(operations) } +/// Extract one LLVM aggregate field into a fresh result of its known type. fn extract_value( result: MlirResult, aggregate: MlirValueUse, @@ -1049,6 +1080,7 @@ fn extract_value( Ok(operation) } +/// Insert a field at its physical LLVM aggregate index. fn insert_value( result: MlirResult, aggregate: MlirValueUse, @@ -1069,20 +1101,24 @@ fn insert_value( 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( @@ -1096,6 +1132,7 @@ fn narrow_vector_type(kind: NarrowLoad) -> MlirType { } } +/// Wrap CUTLASS's narrow floating-point type spelling without widening it. fn narrow_float(spelling: &str) -> MlirType { match spelling { "f16" => MlirType::Float(MlirFloatType::F16), @@ -1105,6 +1142,7 @@ fn narrow_float(spelling: &str) -> MlirType { } } +/// 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], @@ -1112,6 +1150,7 @@ fn vector_type(size: u64, element: MlirType) -> MlirType { } } +/// 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); @@ -1121,6 +1160,7 @@ fn tensor_view(ctx: &Context, value: pliron::value::Value) -> Result Result { let ty = value.get_type(ctx); let ty_ref = ty.deref(ctx); @@ -1130,6 +1170,7 @@ fn scaled_view(ctx: &Context, value: pliron::value::Value) -> Result Result { let ty = value.get_type(ctx); let ty_ref = ty.deref(ctx); @@ -1139,6 +1180,7 @@ fn fragment(ctx: &Context, value: pliron::value::Value) -> Result Result<(), String> { view.verify(ctx).map_err(stringify)?; if view.space != CuteTensorAddressSpaceAttr::Gmem @@ -1160,6 +1202,7 @@ fn require_gemv_tensor_view(ctx: &Context, view: &CuteTensorViewType) -> Result< } } +/// Require exact result/operand counts and no leftover attributes, regions, or successors. fn expect_shape( name: &str, input: &OperationInput, @@ -1181,6 +1224,7 @@ fn expect_shape( Ok(()) } +/// Build a target operation with the supplied operands, results, and source location. fn operation( name: &str, results: Vec, @@ -1194,11 +1238,13 @@ fn operation( 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, @@ -1215,6 +1261,7 @@ fn integer_constant( Ok(operation) } +/// Build a two-operand target operation using the supplied result. fn binary( name: &str, result: MlirResult, @@ -1225,6 +1272,7 @@ fn binary( 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") @@ -1233,10 +1281,12 @@ fn llvm_pointer_type(address_space: u32) -> Result { } } +/// 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_attention.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tests/cooperative.rs similarity index 59% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_attention.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tests/cooperative.rs index a8bfaa0542..74090d89b9 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_attention.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tests/cooperative.rs @@ -3,626 +3,17 @@ * 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. +//! Independent address evaluation for cooperative-copy and matrix-load mappings. -use cute_layout::{ - ComposedLayout, IntTuple, Layout, OffsetUnit, Swizzle, validate_cooperative_copy_plan, +use crate::cute::{ + copy::emit::integer, + mma::ldmatrix::{matrix_address_plan, matrix_offset_fits_i32}, }; -use dialect_cute::{ - attributes::{ - CuteComposedLayoutAttr, CuteCopyAtomAttr, CuteDivisibilityAttr, CuteLayoutAttr, - CuteMatrixRoleAttr, - }, - ops::{CuteCopyG2SOp, CuteLdmatrixOp}, -}; -use dialect_mir::types::MirFP16Type; -use pliron::{ - common_traits::Verify, - context::{Context, Ptr}, - operation::Operation, - r#type::TypeHandle, -}; -use pliron_mlir_export::{ - DropAttribute, MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, - OperationInput, OperationTranslation, TranslationError, TranslationRegistry, - TranslationSession, -}; - -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(()) -} - -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(); - if !elem.deref(ctx).is::() { - return Err("CuTe attention copy supports f16 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, - MlirType::Float(pliron_mlir_export::MlirFloatType::F16), - )?; - 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) - } -} - -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. -fn matrix_offset_fits_i32(layout: &ComposedLayout) -> bool { - let limit = i64::from(i32::MAX); - if layout.unit() != OffsetUnit::Bytes - || layout.inner().modes().len() != 2 - || !layout - .inner() - .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 - .inner() - .shape - .leaves() - .iter() - .zip(&layout.inner().stride.leaves()) - { - 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.outer(); - swizzle.is_valid() - && (0..=limit).contains(&swizzle.y_mask()) - && (0..=limit).contains(&swizzle.z_mask()) -} - -struct MatrixAddressPlan { - low_mask: i64, - 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. -fn matrix_address_plan( - layout: &ComposedLayout, - role: CuteMatrixRoleAttr, -) -> Option { - if !matrix_offset_fits_i32(layout) - || layout.inner().checked_size()? > cute_layout::validation::MAX_STATIC_VALIDATION_ELEMENTS - { - return None; - } - let modes = layout.inner().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.outer(); - 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.inner(), 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.inner(), 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.outer(), low)?; - let local = e.swizzle(layout.outer(), 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) - } -} - -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(()) -} - -fn integer(value: i128, width: u32) -> MlirAttribute { - MlirAttribute::Integer { - value, - ty: MlirType::Integer(width), - } -} - -struct Emitter<'a, 'ctx> { - session: &'a mut TranslationSession<'ctx>, - location: MlirLocation, - operations: Vec, -} -impl<'a, 'ctx> Emitter<'a, 'ctx> { - fn new(session: &'a mut TranslationSession<'ctx>, location: MlirLocation) -> Self { - Self { - session, - location, - operations: vec![], - } - } - fn fresh(&mut self, ty: MlirType) -> (MlirResult, MlirValueUse) { - let id = self.session.fresh_value(); - (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) - } - fn operation(&self, name: &str, operands: Vec) -> Result { - let mut op = MlirOperation::new(name)?; - op.location = self.location.clone(); - op.operands = operands; - Ok(op) - } - 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_) - } - fn c64(&mut self, value: i64) -> Result { - self.constant(value, 64) - } - 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_) - } - fn binary( - &mut self, - name: &str, - a: MlirValueUse, - b: MlirValueUse, - ) -> Result { - self.value(name, a.ty.clone(), vec![a, b]) - } - 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) - } - fn select( - &mut self, - predicate: MlirValueUse, - yes: MlirValueUse, - no: MlirValueUse, - ) -> Result { - self.value("arith.select", yes.ty.clone(), vec![predicate, yes, no]) - } - 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]) - } - } - 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. - fn index( - &mut self, - layout: &Layout, - mut coordinate: MlirValueUse, - ) -> Result { - let sizes = layout.shape.leaves(); - let strides = layout.stride.leaves(); - 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) - } - 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) - } - 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.inner(), row, column)?; - let offset = self.constant(layout.offset(), width)?; - let plain = self.binary("arith.addi", plain, offset)?; - self.swizzle(layout.outer(), plain) - } - 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.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.shift.unsigned_abs()), width)?; - let moved = self.binary( - if swizzle.shift >= 0 { - "arith.shrui" - } else { - "arith.shli" - }, - bits, - shift, - )?; - self.binary("arith.xori", plain, moved) - } -} #[cfg(test)] mod tests { + //! Evaluate exported copy and matrix-load address arithmetic against independent layout expectations. + use crate::{CutlassFullCuteMlir22, MlirConsumerProfile}; use cute_layout::{ComposedLayout, Layout, OffsetUnit, Swizzle}; use dialect_cute::{ @@ -649,6 +40,7 @@ mod tests { }; use pliron_mlir_export::render_module; + /// Construct the canonical swizzled shared tile used by cooperative-copy fixtures. fn shared_layout() -> ComposedLayout { ComposedLayout::new( Swizzle::new(3, 3, 3), @@ -659,14 +51,17 @@ mod tests { .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, @@ -676,6 +71,7 @@ mod tests { 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, @@ -686,8 +82,10 @@ mod tests { 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, @@ -769,6 +167,7 @@ mod tests { 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(); @@ -791,6 +190,7 @@ mod tests { 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 [ @@ -816,6 +216,7 @@ mod tests { } } + /// Find the device function's block through the executable GPU envelope. fn function_block(op: &pliron_mlir_export::MlirOperation) -> &pliron_mlir_export::MlirBlock { if op.name == "func.func" { return &op.regions[0].blocks[0]; @@ -825,6 +226,7 @@ mod tests { .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: &pliron_mlir_export::MlirOperation) -> bool { op.name == "func.func" || op @@ -842,6 +244,7 @@ mod tests { // 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: &pliron_mlir_export::MlirBlock, args: &[u64]) -> Vec> { use pliron_mlir_export::{MlirAttribute as A, MlirFloatType, MlirType as T}; let mut values = block @@ -918,6 +321,7 @@ mod tests { memory } + /// 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(); @@ -978,6 +382,7 @@ mod tests { } } + /// 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 @@ -1122,6 +527,7 @@ mod tests { } } + /// 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 [ @@ -1154,6 +560,7 @@ mod tests { } } + /// 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(); @@ -1197,6 +604,7 @@ mod tests { ); } + /// 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 [ @@ -1235,6 +643,7 @@ mod tests { } } + /// 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); @@ -1290,6 +699,7 @@ mod tests { 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; @@ -1490,6 +900,7 @@ mod tests { } } + /// 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 inner in ["(1024,2048):(4096,2)", "(8,64):(128,2)", "(16,8):(16,2)"] { @@ -1508,6 +919,7 @@ mod tests { } } + /// 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(); 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..23ea157036 --- /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 dialect_cute::attributes::CuteEpilogueSyncPhaseAttr; +use pliron_mlir_export::{MlirAttribute, MlirLocation, MlirType, MlirValueUse}; + +#[cfg(test)] +mod tests { + //! Focused checks for GEMM byte offsets, packed lanes, and synchronization ordering. + + use super::*; + use pliron_mlir_export::{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/cute_gemm.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_gemm.rs deleted file mode 100644 index 0e6e656a08..0000000000 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute_gemm.rs +++ /dev/null @@ -1,3118 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! SM120 block-scaled GEMM mappings for the official CUTLASS 4.7 compiler. -//! -//! The source dialect keeps scheduler, pipeline, TMA, MMA, and epilogue -//! operations visible. This pack removes compile-time handles and maps their -//! runtime work to operations accepted by the public PRE_COMPILED pipeline -//! in `libCutlassCompiler.so`. - -use dialect_cute::{ - attributes::{ - CuteCountedCtaBarrierAttr, CuteEpilogueHalfAttr, CuteEpilogueSyncPhaseAttr, - CuteMmaCarrierKindAttr, CutePipelineRoleAttr, CutePipelineStateAttr, - CuteTensorAddressSpaceAttr, 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}, - types::{MirArrayType, MirStructType, MirTupleType}, -}; -use pliron::{ - common_traits::Verify, - context::{Context, Ptr}, - operation::Operation, - r#type::{TypeHandle, Typed}, -}; -use pliron_mlir_export::{ - DropAttribute, MlirAttribute, MlirBlock, MlirLocation, MlirOperation, MlirRegion, MlirResult, - MlirType, MlirValueUse, OperationInput, OperationTranslation, TranslationError, - TranslationRegistry, TranslationSession, TypeTranslation, -}; - -const PIPELINE_WAIT_TICKS: i128 = 10_000_000; -const LDMATRIX_BYTE_SWIZZLE_MASK: u64 = 0x180; -const LDMATRIX_BYTE_SWIZZLE_SHIFT: u64 = 3; -const COUNTED_CTA_SYNC_OP: &str = "nvvm.barrier.cta.sync"; - -/// 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) -} - -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)) - } -} - -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() - } -} - -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(), - } - } -} - -struct StaticContractTypeTranslation; - -impl TypeTranslation for StaticContractTypeTranslation { - fn translate( - &self, - _ctx: &Context, - _source: TypeHandle, - _registry: &TranslationRegistry, - ) -> Result { - Ok(MlirType::Tuple(vec![])) - } -} - -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; - let mut target = operation( - "llvm.insertvalue", - input.results, - input.operands, - &input.location, - )?; - target - .properties - .insert("position".into(), MlirAttribute::DenseI64Array(vec![index])); - Ok(vec![target]) - } -} - -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]) - } -} - -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) - } -} - -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]) - } -} - -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) - } -} - -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) - } -} - -struct IdentityTranslation(&'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, - )?]) - } -} - -struct IdentityFirstOperandTranslation(&'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. - -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) - } -} - -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, - )?, - ]) - } -} - -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, - )?]) - } -} - -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) - } -} - -struct PipelineWaitTranslation { - empty_ring: bool, - 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) - } -} - -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) - } -} - -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) - } -} - -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) - } -} - -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) - } -} - -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) - } -} - -struct TmaStoreWaitTranslation { - 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]) - } -} - -struct NoOperandRenameTranslation(&'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)?]) - } -} - -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) - } -} - -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) - } -} - -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(()) -} - -struct StaticMultiIdentityTranslation { - name: &'static str, - 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() - } -} - -#[derive(Clone)] -struct CarrierLayout { - target: MlirType, - kind: CarrierKind, -} - -#[derive(Clone)] -enum CarrierKind { - Leaf, - Aggregate(Vec<(i64, CarrierLayout)>), -} - -impl CarrierLayout { - 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), - }) - } - - fn leaf_count(&self) -> usize { - match &self.kind { - CarrierKind::Leaf => 1, - CarrierKind::Aggregate(children) => { - children.iter().map(|(_, child)| child.leaf_count()).sum() - } - } - } -} - -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(()) - } - } -} - -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(()) -} - -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) -} - -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) - } -} - -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) - } -} - -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) - } -} - -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) - } -} - -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()) -} - -struct TiledGemmTranslation; - -impl OperationTranslation for TiledGemmTranslation { - fn translate( - &self, - ctx: &Context, - source: Ptr, - mut input: OperationInput, - session: &mut TranslationSession<'_>, - ) -> Result, String> { - CuteTiledGemmOp::wrap(source) - .verify(ctx) - .map_err(stringify)?; - 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(), - &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) - } -} - -#[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, - 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 = cast_int( - session, - operations, - scale_a, - MlirType::Integer(16), - "arith.trunci", - location, - )?; - let scale_b = 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(pliron_mlir_export::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(pliron_mlir_export::MlirFloatType::Other( - "f4E2M1FN".into(), - ))), - ); - operation_.properties.insert( - "b_type".into(), - MlirAttribute::Type(MlirType::Float(pliron_mlir_export::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(pliron_mlir_export::MlirFloatType::Other( - "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: 32, - ty: MlirType::Integer(32), - }, - ); - operations.push(operation_); - uses.try_into() - .map_err(|_| "SM120 MXF4 MMA did not produce four accumulators".to_owned()) -} - -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)?; - - 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) - } -} - -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(pliron_mlir_export::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) -} - -fn f16x2_intrinsic_operands(low: MlirValueUse, high: MlirValueUse) -> [MlirValueUse; 2] { - [high, low] -} - -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_) -} - -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_) -} - -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()) -} - -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.inner().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)) -} - -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_) -} - -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, - ) -} - -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, - ) -} - -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])) -} - -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_) -} - -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), - ); -} - -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_) -} - -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_) -} - -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_)) -} - -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_) -} - -fn binary( - name: &str, - result: MlirResult, - left: MlirValueUse, - right: MlirValueUse, - location: &MlirLocation, -) -> Result { - operation(name, vec![result], vec![left, right], location) -} - -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) -} - -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(()) -} - -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) -} - -fn mbarrier_init( - barrier: MlirValueUse, - arrival_count: MlirValueUse, - location: &MlirLocation, -) -> Result { - operation( - "nvvm.mbarrier.init", - vec![], - vec![barrier, arrival_count], - location, - ) -} - -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, - ) -} - -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) -} - -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) -} - -fn epilogue_gep_element_type() -> MlirType { - MlirType::Float(pliron_mlir_export::MlirFloatType::F16) -} - -#[cfg(test)] -const fn ldmatrix_swizzled_byte_offset(plain: u64) -> u64 { - plain ^ ((plain & LDMATRIX_BYTE_SWIZZLE_MASK) >> LDMATRIX_BYTE_SWIZZLE_SHIFT) -} - -fn fresh(session: &mut TranslationSession<'_>, ty: MlirType) -> (MlirResult, MlirValueUse) { - let id = session.fresh_value(); - (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) -} - -fn generic_pointer_type() -> Result { - MlirType::dialect("!llvm.ptr") -} - -fn shared_pointer_type() -> Result { - MlirType::dialect("!llvm.ptr<3>") -} - -fn u64_bits(value: u64) -> i128 { - i64::from_ne_bytes(value.to_ne_bytes()) as i128 -} - -fn stringify(error: impl std::fmt::Display) -> String { - error.to_string() -} - -#[cfg(test)] -mod tests { - use super::*; - use pliron_mlir_export::{MlirFloatType, MlirValueId}; - - fn value(id: u64, ty: MlirType) -> MlirValueUse { - MlirValueUse { - id: MlirValueId(id), - ty, - } - } - - #[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); - } - - #[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); - } - - #[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) - ); - } - - #[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 index 87af4b7102..c2d12a4614 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/grid_constant.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/grid_constant.rs @@ -19,12 +19,16 @@ 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. @@ -140,7 +144,7 @@ pub(crate) fn take_parameters( "grid-constant parameter {index} pointee metadata differs from its pointer type" )); } - let size = crate::mir_memory::source_stored_size(ctx, pointee) + 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") diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/lib.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/lib.rs index 9eecaf9988..ebe6137758 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/lib.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/lib.rs @@ -3,46 +3,48 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! CUDA Oxide mappings for textual MLIR export. +//! cuda-oxide mappings for textual MLIR export. //! -//! The generic exporter builds typed MLIR text. This crate maps CUDA Oxide +//! The generic exporter builds typed MLIR text. This crate maps cuda-oxide //! operations to the types and operations accepted by the pinned CUTLASS compiler. +//! +//! 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 = pliron_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 cute_attention; -mod cute_gemm; -mod cute_gemv; -mod cute_sm100; -mod cute_sm100_attention; mod grid_constant; -mod mir_checked; -mod mir_core; -mod mir_memory; -mod mir_shared; -mod nvvm_attention; -mod nvvm_cluster; -mod nvvm_numeric; -mod nvvm_sm120; -mod nvvm_sregs; -mod nvvm_sync_restrict; -mod nvvm_tcgen05; +mod mir; +mod nvvm; mod packs; mod primitives; mod profile; pub use cute::register_cute_elementwise_pack; -pub(crate) use cute_attention::register_cute_attention_pack; -pub(crate) use cute_gemm::register_cute_gemm_pack; -pub(crate) use cute_gemv::register_cute_cutlass_profile_pack; -pub(crate) use cute_sm100::register_cute_sm100_pack; -pub use mir_core::register_mir_core_pack; -pub use nvvm_attention::register_nvvm_attention_pack; -pub use nvvm_cluster::register_nvvm_cluster_pack; -pub use nvvm_numeric::register_nvvm_numeric_pack; -pub use nvvm_sm120::register_nvvm_sm120_pack; -pub use nvvm_sregs::register_nvvm_sreg_pack; -pub use nvvm_sync_restrict::register_nvvm_sync_restrict_pack; -pub use nvvm_tcgen05::register_nvvm_tcgen05_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 similarity index 91% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/mir_checked.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/checked.rs index 1ec33f38dd..01775453ab 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir_checked.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/checked.rs @@ -25,6 +25,7 @@ use pliron_mlir_export::{ TranslationSession, TypeTranslation, }; +/// 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)?; @@ -37,6 +38,7 @@ pub(crate) fn register(r: &mut TranslationRegistry) -> Result<(), TranslationErr Ok(()) } +/// Validate Direct/Single layout and build byte storage with the original ABI alignment. pub(crate) fn enum_storage_type(e: &MirEnumType) -> Result { if !matches!( e.layout_kind, @@ -77,9 +79,11 @@ pub(crate) fn enum_storage_type(e: &MirEnumType) -> Result { 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. fn check_payload_storage(ctx: &Context, ty: TypeHandle) -> Result<(), String> { let t = ty.deref(ctx); if t.downcast_ref::() @@ -128,6 +132,7 @@ impl TypeTranslation for EnumType { } } +/// Build an operation with the supplied result IDs, operands, and source location. fn op( name: &str, operands: Vec, @@ -140,10 +145,12 @@ fn op( 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, @@ -161,7 +168,9 @@ fn integer( } // 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. @@ -172,10 +181,12 @@ const ENUM_SCRATCH_SLOT: &str = "cuda_oxide.enum_scratch_slot"; /// 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 { @@ -252,6 +263,7 @@ pub(crate) fn hoist_enum_scratch(function: &mut MlirOperation) -> Result<(), Str Ok(()) } +/// Allocate one private, aligned enum scratch slot and mark it for entry-block hoisting. fn slot( ty: MlirType, align: u64, @@ -280,6 +292,7 @@ fn slot( out.push(a); Ok(p) } +/// Offset an enum scratch pointer by its rustc-provided byte position. fn byte_pointer( base: MlirValueUse, offset: u64, @@ -308,6 +321,7 @@ fn byte_pointer( out.push(gep); Ok(p) } +/// Write one enum tag or payload to potentially unaligned byte storage. fn store( value: MlirValueUse, pointer: MlirValueUse, @@ -326,6 +340,7 @@ fn store( out.push(st); Ok(()) } +/// Read one enum tag or payload from potentially unaligned byte storage. fn load( pointer: MlirValueUse, result: MlirResult, @@ -343,9 +358,13 @@ fn load( 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 { @@ -432,7 +451,7 @@ impl OperationTranslation for EnumOp { .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]) + 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()); @@ -477,7 +496,7 @@ impl OperationTranslation for EnumOp { .map(|v| *v as usize) .sum::() + field; - if crate::mir_memory::source_stored_size(ctx, e.all_field_types[flat]) + 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()); @@ -495,7 +514,11 @@ impl OperationTranslation for EnumOp { } } -struct Checked(&'static str); +/// 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, @@ -579,10 +602,13 @@ impl OperationTranslation for Checked { fields.push(field); } input.operands = fields; - out.extend(crate::mir_memory::ConstructStructTranslation.translate(ctx, source, input, s)?); + 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( @@ -611,8 +637,11 @@ impl OperationTranslation for Unreachable { #[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(); @@ -638,6 +667,7 @@ mod tests { .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(); @@ -659,6 +689,7 @@ mod tests { 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(); @@ -692,6 +723,7 @@ mod tests { 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}; @@ -709,6 +741,7 @@ mod tests { 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, @@ -719,6 +752,7 @@ mod tests { 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, @@ -741,6 +775,7 @@ mod tests { 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); @@ -887,6 +922,7 @@ mod tests { 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 pliron_mlir_export::{MlirBlock, MlirBlockId, MlirRegion, MlirValueId}; let mut count = MlirOperation::new("arith.constant").unwrap(); @@ -938,6 +974,7 @@ mod tests { 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(); @@ -950,6 +987,7 @@ mod tests { 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 { 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 similarity index 92% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/mir_core.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/core.rs index d3afd1a46d..4d64b2eea7 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir_core.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/core.rs @@ -80,12 +80,13 @@ pub fn register_mir_core_pack(registry: &mut TranslationRegistry) -> Result<(), 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)?; + 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 { @@ -105,7 +106,8 @@ impl AttributeTranslation for MirFp16AttributeTranslation { } } -struct RenameOperation(&'static str); +/// Forward translated structure to a standard MLIR operation spelling. +struct RenameOperation(#[doc = "Target MLIR operation spelling."] &'static str); impl OperationTranslation for RenameOperation { fn translate( @@ -119,6 +121,7 @@ impl OperationTranslation for RenameOperation { } } +/// Map MIR functions, launch contracts, slice ABI, and enum scratch to executable MLIR. struct FunctionTranslation; impl OperationTranslation for FunctionTranslation { @@ -155,7 +158,7 @@ impl OperationTranslation for FunctionTranslation { } } } - crate::mir_checked::hoist_enum_scratch(&mut target)?; + crate::mir::checked::hoist_enum_scratch(&mut target)?; Ok(vec![target]) } } @@ -194,7 +197,7 @@ fn lower_alwaysinline_contract(target: &mut MlirOperation) -> Result<(), String> } } -/// Move CUDA Oxide's three scalar block-contract attributes to the one array +/// 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> { @@ -260,6 +263,7 @@ fn move_cluster_contract(target: &mut MlirOperation) -> Result<(), String> { } } +/// Remove a named launch attribute and require a positive i32-representable value. fn take_positive_i32_attribute( target: &mut MlirOperation, name: &str, @@ -473,6 +477,7 @@ fn flatten_kernel_slice_abi( Ok(()) } +/// Rebuild conditional branch successor operand segments for cf.cond_br. struct ConditionalBranchTranslation; impl OperationTranslation for ConditionalBranchTranslation { @@ -504,8 +509,11 @@ impl OperationTranslation for ConditionalBranchTranslation { } } +/// 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, } @@ -543,11 +551,17 @@ impl OperationTranslation for ConstantTranslation { } } +/// 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, } @@ -578,7 +592,8 @@ impl OperationTranslation for NumericBinaryTranslation { } } -struct IntegerBinaryTranslation(&'static str); +/// 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( @@ -593,7 +608,9 @@ impl OperationTranslation for IntegerBinaryTranslation { } } +/// 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, } @@ -675,6 +692,7 @@ impl OperationTranslation for ShiftTranslation { } } +/// Negate using integer zero subtraction or floating-point negation as appropriate. struct NegTranslation; impl OperationTranslation for NegTranslation { @@ -723,6 +741,7 @@ impl OperationTranslation for NegTranslation { } } +/// Complement every bit of the source integer width. struct NotTranslation; impl OperationTranslation for NotTranslation { @@ -770,12 +789,19 @@ impl OperationTranslation for NotTranslation { } } +/// 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, } @@ -830,6 +856,7 @@ impl OperationTranslation for ComparisonTranslation { } } +/// Translate scalar, bitwise, and pointer conversions without losing source address semantics. struct CastTranslation; impl OperationTranslation for CastTranslation { @@ -993,6 +1020,7 @@ impl OperationTranslation for CastTranslation { } } +/// Recover a source MIR pointer's numeric address space before opaque-pointer translation. fn source_pointer_address_space( ctx: &Context, source: Ptr, @@ -1010,13 +1038,25 @@ fn source_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 { - Integer { width: u32, signed: bool }, - Float { width: u32 }, + /// 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, @@ -1024,6 +1064,7 @@ impl Scalar { } } +/// Read scalar classification from a selected source operand. fn source_operand_scalar( ctx: &Context, source: Ptr, @@ -1033,6 +1074,7 @@ fn source_operand_scalar( scalar(ctx, ty) } +/// Read scalar classification from a selected source result. fn source_result_scalar( ctx: &Context, source: Ptr, @@ -1042,6 +1084,7 @@ fn source_result_scalar( 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::() { @@ -1067,6 +1110,7 @@ fn scalar(ctx: &Context, ty: TypeHandle) -> Result { Ok(Scalar::Float { width }) } +/// Require an integer source value and retain its bit width and signedness. fn source_integer( ctx: &Context, source: Ptr, @@ -1080,6 +1124,7 @@ fn source_integer( .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; @@ -1091,6 +1136,7 @@ fn renamed(name: &str, input: OperationInput) -> Result { Ok(target) } +/// Remove an optional source attribute and reattach it under the target property name. fn move_property( operation: &mut MlirOperation, source_name: &str, @@ -1106,8 +1152,11 @@ fn move_property( #[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 pliron_mlir_export::{MlirAttribute, MlirOperation, MlirType}; @@ -1169,12 +1218,14 @@ mod tests { value::Value, }; + /// 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, @@ -1198,6 +1249,7 @@ mod tests { (operation, region) } + /// Construct a source operation of the requested concrete kind and append it to the block. fn operation( ctx: &mut Context, block: Ptr, @@ -1210,6 +1262,7 @@ mod tests { 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(); @@ -1219,10 +1272,12 @@ mod tests { ) } + /// 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}; @@ -1281,6 +1336,7 @@ mod tests { } } + /// 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(); @@ -1400,6 +1456,7 @@ mod tests { assert_eq!(text, 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(); @@ -1428,6 +1485,7 @@ mod tests { ); } + /// 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(); @@ -1457,6 +1515,7 @@ mod tests { ); } + /// 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)] @@ -1502,6 +1561,7 @@ mod tests { assert!(!text.contains("llvm.ptrtoint"), "{text}"); } + /// Check explicit inline(always) survives lowering for a scalar-only helper. #[test] fn alwaysinline_pointer_free_helper_uses_llvm_passthrough() { let mut ctx = context(); @@ -1520,6 +1580,7 @@ mod tests { 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(); @@ -1541,6 +1602,7 @@ mod tests { assert!(!text.contains(r#"alwaysinline = "true""#), "{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(); 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 similarity index 93% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/mir_memory.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/memory.rs index 7ff9ddb4e6..9301122744 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir_memory.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/memory.rs @@ -35,6 +35,7 @@ use pliron_mlir_export::{ TranslationSession, TypeTranslation, }; +/// Register reviewed pointer, physical aggregate, call, and memory mappings. pub(crate) fn register_mir_memory_pack( registry: &mut TranslationRegistry, ) -> Result<(), TranslationError> { @@ -69,6 +70,7 @@ pub(crate) fn register_mir_memory_pack( Ok(()) } +/// Erase a MIR pointee while preserving its LLVM address space. struct PointerTypeTranslation; impl TypeTranslation for PointerTypeTranslation { @@ -86,6 +88,7 @@ impl TypeTranslation for PointerTypeTranslation { } } +/// Represent a Rust slice as a data pointer plus 64-bit element count. struct SliceTypeTranslation; impl TypeTranslation for SliceTypeTranslation { @@ -99,6 +102,7 @@ impl TypeTranslation for SliceTypeTranslation { } } +/// Preserve the disjoint slice's runtime pointer and extent carrier. struct DisjointSliceTypeTranslation; impl TypeTranslation for DisjointSliceTypeTranslation { @@ -120,6 +124,7 @@ impl TypeTranslation for DisjointSliceTypeTranslation { } } +/// Translate fixed-length array storage with its element type and length. struct ArrayTypeTranslation; impl TypeTranslation for ArrayTypeTranslation { @@ -140,6 +145,7 @@ impl TypeTranslation for ArrayTypeTranslation { } } +/// Build an LLVM struct with verified Rust field order, offsets, and padding. struct StructTypeTranslation; impl TypeTranslation for StructTypeTranslation { @@ -153,6 +159,7 @@ impl TypeTranslation for StructTypeTranslation { } } +/// Build an LLVM struct retaining the tuple's physical Rust layout. struct TupleTypeTranslation; impl TypeTranslation for TupleTypeTranslation { @@ -166,6 +173,7 @@ impl TypeTranslation for TupleTypeTranslation { } } +/// Allocate local storage with source element type, count, and explicit alignment. struct AllocaTranslation; impl OperationTranslation for AllocaTranslation { @@ -221,6 +229,7 @@ impl OperationTranslation for AllocaTranslation { } } +/// Load a verified pointee while preserving the source memory access contract. struct LoadTranslation; impl OperationTranslation for LoadTranslation { @@ -246,6 +255,7 @@ impl OperationTranslation for LoadTranslation { } } +/// Store a value to its verified pointee with explicit alignment and memory ordering. struct StoreTranslation; impl OperationTranslation for StoreTranslation { @@ -278,6 +288,7 @@ impl OperationTranslation for StoreTranslation { } } +/// Compute a pointer offset in source pointee elements. struct PointerOffsetTranslation; impl OperationTranslation for PointerOffsetTranslation { @@ -306,6 +317,7 @@ impl OperationTranslation for PointerOffsetTranslation { } } +/// Extract a logical field using its physical LLVM aggregate slot. struct ExtractFieldTranslation; impl OperationTranslation for ExtractFieldTranslation { @@ -370,8 +382,10 @@ impl OperationTranslation for ExtractFieldTranslation { } } +/// 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 { @@ -505,6 +519,7 @@ impl OperationTranslation for ConstructStructTranslation { } } +/// Compute a logical field address through the verified physical aggregate layout. struct FieldAddressTranslation; impl OperationTranslation for FieldAddressTranslation { @@ -548,6 +563,7 @@ impl OperationTranslation for FieldAddressTranslation { } } +/// Compute the address of a dynamic element within a fixed array. struct ArrayElementAddressTranslation; impl OperationTranslation for ArrayElementAddressTranslation { @@ -572,6 +588,7 @@ impl OperationTranslation for ArrayElementAddressTranslation { } } +/// Extract an array element while preserving constant and dynamic index behavior. struct ExtractArrayElementTranslation; impl OperationTranslation for ExtractArrayElementTranslation { @@ -688,6 +705,7 @@ impl OperationTranslation for ExtractArrayElementTranslation { } } +/// Map reviewed Rust intrinsics directly and preserve ordinary function call signatures. struct CallTranslation; impl OperationTranslation for CallTranslation { @@ -802,6 +820,7 @@ impl OperationTranslation for CallTranslation { } } +/// Keep a MIR assertion as an effectful cf.assert before following control flow. struct AssertTranslation; impl OperationTranslation for AssertTranslation { @@ -835,6 +854,7 @@ impl OperationTranslation for AssertTranslation { } } +/// Erase a source-only operation after checking that it has no runtime results. struct DropOperation; impl OperationTranslation for DropOperation { @@ -857,13 +877,18 @@ impl OperationTranslation for DropOperation { } } +/// 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, @@ -974,6 +999,7 @@ pub(crate) fn aggregate_plan( }) } +/// 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!( @@ -993,10 +1019,11 @@ fn validate_memory_order(order: &[usize], fields: usize) -> Result<(), String> { 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(e) = source_ref.downcast_ref::() { - return e.total_size == 0 && crate::mir_checked::enum_storage_type(e).is_ok(); + return e.total_size == 0 && crate::mir::checked::enum_storage_type(e).is_ok(); } if source_ref.downcast_ref::().is_some() { return true; @@ -1019,10 +1046,11 @@ fn source_is_zero_sized(ctx: &Context, source: TypeHandle) -> bool { 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(e) = source_ref.downcast_ref::() { - return crate::mir_checked::enum_storage_type(e) + return crate::mir::checked::enum_storage_type(e) .ok() .map(|_| e.total_size); } @@ -1078,10 +1106,11 @@ pub(crate) fn source_stored_size(ctx: &Context, source: TypeHandle) -> Option Option { let source_ref = source.deref(ctx); if let Some(e) = source_ref.downcast_ref::() { - return crate::mir_checked::enum_storage_type(e) + return crate::mir::checked::enum_storage_type(e) .ok() .map(|_| e.abi_align); } @@ -1135,6 +1164,7 @@ pub(crate) fn source_natural_alignment(ctx: &Context, source: TypeHandle) -> Opt None } +/// Read explicit Rust aggregate or enum alignment, falling back to natural alignment. fn source_abi_align(ctx: &Context, source: TypeHandle) -> Option { let source_ref = source.deref(ctx); if let Some(structure) = source_ref.downcast_ref::() @@ -1150,6 +1180,7 @@ fn source_abi_align(ctx: &Context, source: TypeHandle) -> Option { 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; @@ -1165,6 +1196,7 @@ fn natural_aggregate_size(ctx: &Context, fields: &[TypeHandle], order: &[usize]) 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}")); @@ -1175,6 +1207,7 @@ fn align_up(value: u64, alignment: u64) -> Result { .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 @@ -1183,6 +1216,7 @@ fn pointer_pointee(ctx: &Context, pointer: TypeHandle) -> Result Result { if address_space == 0 { MlirType::dialect("!llvm.ptr") @@ -1191,10 +1225,12 @@ fn llvm_pointer_type(address_space: u32) -> Result { } } +/// 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() @@ -1204,6 +1240,7 @@ fn llvm_struct_type(fields: &[MlirType]) -> Result { 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(), @@ -1246,6 +1283,7 @@ pub(crate) fn type_spelling(ty: &MlirType) -> String { } } +/// 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()), @@ -1253,6 +1291,7 @@ fn only_result(name: &str, input: &OperationInput) -> Result } } +/// Reject source attributes not consumed by the recipe. fn ensure_empty_attributes(name: &str, input: &OperationInput) -> Result<(), String> { if input.attributes.is_empty() { Ok(()) @@ -1264,6 +1303,7 @@ fn ensure_empty_attributes(name: &str, input: &OperationInput) -> Result<(), Str } } +/// 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; @@ -1275,6 +1315,7 @@ fn renamed(name: &str, input: OperationInput) -> Result { 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(), @@ -1285,6 +1326,7 @@ fn memory_ordering(operation: &mut MlirOperation) { ); } +/// 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( @@ -1297,6 +1339,7 @@ fn add_alignment(operation: &mut MlirOperation, alignment: Option) { } } +/// Set GEP element and index metadata without promising no-wrap behavior. fn gep_properties( operation: &mut MlirOperation, element_type: MlirType, @@ -1321,6 +1364,8 @@ fn gep_properties( #[cfg(test)] mod tests { + //! Physical aggregate, dynamic array, and effectful assertion translation checks. + use crate::{ CutlassFullCuteMlir22, MlirConsumerProfile, profile::render_mapping_module_without_cutlass_envelope, @@ -1350,6 +1395,7 @@ mod tests { value::Value, }; + /// Construct a typed source operation and append it to the test block. fn operation( ctx: &mut Context, block: Ptr, @@ -1362,10 +1408,12 @@ mod tests { 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) } + /// Check assertion lowering preserves its execution position before later memory/control-flow effects. #[test] fn assertion_remains_an_effectful_check_before_following_operations() { let mut ctx = Context::new(); @@ -1407,6 +1455,7 @@ mod tests { assert!(!text.contains("\"cf.br\""), "{text}"); } + /// 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] { @@ -1459,6 +1508,7 @@ mod tests { } } + /// 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(); 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..d3291b90a0 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/mod.rs @@ -0,0 +1,11 @@ +/* + * 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; 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 similarity index 95% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/mir_shared.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/shared.rs index 69559a50a4..4fc2774f8d 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir_shared.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/shared.rs @@ -17,9 +17,11 @@ use pliron_mlir_export::{ 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( @@ -60,7 +62,7 @@ impl OperationTranslation for SharedAllocation { let alignment = allocation .get_alignment_value(ctx) .filter(|align| *align != 0) - .or_else(|| crate::mir_memory::source_natural_alignment(ctx, element)) + .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()); @@ -73,7 +75,7 @@ impl OperationTranslation for SharedAllocation { 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) + crate::mir::memory::type_spelling(&ty) ))?; let mut global = MlirOperation::new("llvm.mlir.global")?; // Preserve source identity while deduplicating: signedness and enum @@ -138,7 +140,9 @@ impl OperationTranslation for SharedAllocation { } } +/// 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, @@ -215,6 +219,7 @@ pub(crate) fn register_dynamic_shared( ) -> 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( 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 similarity index 93% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_attention.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/attention.rs index fa174df3a5..c0d68a95be 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_attention.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/attention.rs @@ -8,7 +8,7 @@ //! Selectors and canonical intrinsic identities are checked before dropping //! source attributes. Only the exact variants used by attention are supported. -use crate::nvvm_sm120::{constant_integer, fresh, operation, property, require_input}; +use crate::nvvm::sm120::{constant_integer, fresh, operation, property, require_input}; use dialect_nvvm::ops::*; use pliron::{ common_traits::Verify, @@ -33,6 +33,7 @@ pub fn register_nvvm_attention_pack( 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> { @@ -41,6 +42,7 @@ pub(crate) fn register_nvvm_async_control_pack( Ok(()) } +/// Map reviewed attention math while preserving approximation, subnormal mode, and identity. struct MathTranslation; impl OperationTranslation for MathTranslation { fn translate( @@ -57,7 +59,7 @@ impl OperationTranslation for MathTranslation { if math.get_attr_nvvm_scalar_math_operation(ctx).as_deref() == Some(&ScalarMathOperationAttr::Lg2) { - return crate::nvvm_numeric::MathTranslation.translate(ctx, source, input, session); + return crate::nvvm::numeric::MathTranslation.translate(ctx, source, input, session); } let selectors = ( math.get_attr_nvvm_scalar_math_format(ctx), @@ -119,6 +121,7 @@ impl OperationTranslation for MathTranslation { } } +/// Preserve butterfly-shuffle membership, lane offset, and warp clamp. struct ShuffleTranslation; impl OperationTranslation for ShuffleTranslation { fn translate( @@ -153,6 +156,7 @@ impl OperationTranslation for ShuffleTranslation { } } +/// Commit preceding asynchronous copies into a hardware group. struct CommitTranslation; impl OperationTranslation for CommitTranslation { fn translate( @@ -175,6 +179,7 @@ impl OperationTranslation for CommitTranslation { } } +/// Wait for a legal compile-time number of outstanding asynchronous-copy groups. struct WaitTranslation; impl OperationTranslation for WaitTranslation { fn translate( 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 similarity index 87% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_cluster.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/cluster.rs index 2e45f7e71d..fca26ed6b7 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_cluster.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/cluster.rs @@ -27,6 +27,7 @@ use pliron_mlir_export::{ 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 { @@ -231,29 +232,70 @@ pub fn register_nvvm_cluster_pack( Ok(()) } +/// Select one reviewed cluster instruction contract and its structural selectors. enum Kind { - Sreg(&'static str), - NoOperands(&'static str), + /// 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, - CountedBarrier { arrive: bool, aligned: bool }, + /// 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, - ExpectTx { cluster: bool }, - WaitParity { cluster: bool, relaxed: bool }, - RemoteArrive { relaxed: bool }, + /// 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, - WaitGroup { read: bool }, + /// 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>, } @@ -601,6 +643,7 @@ impl OperationTranslation for Recipe { } } +/// Preserve source cluster-barrier mode and native intrinsic identity. struct ClusterBarrierTranslation; impl OperationTranslation for ClusterBarrierTranslation { @@ -637,6 +680,7 @@ impl OperationTranslation for ClusterBarrierTranslation { } } +/// 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())) { @@ -652,13 +696,17 @@ fn check_marker(input: &mut OperationInput, expected: &str) -> Result<(), String Ok(()) } +/// Operand/result type constraint used by explicit NVVM signatures. #[derive(Clone, Copy, Debug)] enum TypeSpec { - Int(u32), - Pointer(&'static [u32]), + /// 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), @@ -669,6 +717,7 @@ impl TypeSpec { } } +/// Check operand/result counts and each translated type against the reviewed signature. fn signature( input: &OperationInput, operands: &[TypeSpec], @@ -702,6 +751,7 @@ fn signature( 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")?; @@ -711,6 +761,7 @@ fn constant_group(ctx: &Context, source: Ptr) -> Result 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 @@ -729,6 +780,7 @@ fn constant_integer(ctx: &Context, source: Ptr, index: usize) -> Resu .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() @@ -737,6 +789,7 @@ fn pointer_type(space: u32) -> Result { }) } +/// Build a target operation with the supplied operands, results, and source location. fn operation( name: &str, results: Vec, @@ -750,17 +803,20 @@ fn operation( 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, @@ -774,6 +830,7 @@ fn cast( Ok(value) } +/// Convert pointer space explicitly while retaining any required integer-address bridge. fn pointer_cast( session: &mut TranslationSession<'_>, operations: &mut Vec, @@ -798,6 +855,8 @@ fn pointer_cast( #[cfg(test)] mod tests { + //! Contract and rejection tests for native cluster, barrier, and TMA mappings. + use super::*; use crate::{ CutlassFullCuteMlir22, MlirConsumerProfile, @@ -822,16 +881,24 @@ mod tests { }; 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); @@ -878,10 +945,12 @@ mod tests { } } + /// 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, @@ -904,6 +973,7 @@ mod tests { 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 }; @@ -926,6 +996,7 @@ mod tests { 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, @@ -946,6 +1017,7 @@ mod tests { } } + /// 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(); @@ -964,6 +1036,7 @@ mod tests { } } + /// 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] { @@ -994,6 +1067,7 @@ mod tests { 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(); @@ -1023,6 +1097,7 @@ mod tests { } } + /// 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 [ @@ -1057,6 +1132,7 @@ mod tests { ); } + /// 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(); @@ -1084,6 +1160,7 @@ mod tests { 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(); @@ -1107,6 +1184,7 @@ mod tests { 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(); @@ -1118,6 +1196,7 @@ mod tests { 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(); @@ -1133,6 +1212,7 @@ mod tests { 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(); @@ -1177,6 +1257,7 @@ mod tests { ); } + /// 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)] { @@ -1211,6 +1292,7 @@ mod tests { } } + /// 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] { 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 similarity index 88% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_numeric.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/numeric.rs index 8d425d0fdf..839804fe1b 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_numeric.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/numeric.rs @@ -4,7 +4,7 @@ */ //! 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::nvvm::tcgen05::{fresh, operation, require_input}; use dialect_nvvm::ops::*; use pliron::{ common_traits::Verify, @@ -16,8 +16,10 @@ use pliron_mlir_export::{ OperationTranslation, TranslationError, TranslationRegistry, TranslationSession, }; +/// 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> { @@ -51,10 +53,11 @@ pub(crate) fn register_nvvm_profile_numeric_pack( crate::register_nvvm_attention_pack(registry) } else { register_nvvm_numeric_pack(registry)?; - crate::nvvm_attention::register_nvvm_async_control_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> { @@ -65,6 +68,7 @@ fn register_nvvm_numeric_common_pack( 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 { @@ -79,6 +83,7 @@ fn register_nvvm_numeric_common_pack( 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 { @@ -135,6 +140,7 @@ fn register_nvvm_numeric_common_pack( 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( @@ -156,11 +162,17 @@ fn intrinsic(input: OperationInput, name: &str) -> Result 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 { @@ -194,6 +206,7 @@ impl OperationTranslation for IntrinsicTranslation { } } +/// Map reviewed scalar math to exact native intrinsics with preserved precision and subnormal mode. pub(crate) struct MathTranslation; impl OperationTranslation for MathTranslation { fn translate( @@ -242,6 +255,7 @@ impl OperationTranslation for MathTranslation { } } +/// Preserve scalar arithmetic format, rounding, subnormal, and saturation selectors. struct ArithmeticTranslation; impl OperationTranslation for ArithmeticTranslation { fn translate( @@ -280,9 +294,13 @@ impl OperationTranslation for ArithmeticTranslation { /// 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 { @@ -331,7 +349,8 @@ impl OperationTranslation for PackedArithmeticTranslation { } } -struct RegisterTranslation(bool); +/// 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, @@ -376,6 +395,7 @@ impl OperationTranslation for RegisterTranslation { Ok(vec![op]) } } +/// Convert low-first FP32 operands into the exact packed FP16 register representation. struct PackTranslation; impl OperationTranslation for PackTranslation { fn translate( @@ -424,7 +444,10 @@ impl OperationTranslation for PackTranslation { } /// Native MLIR operations retain the compiler's PTX lowering for these instructions. -struct SyncTranslation(u8); +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, @@ -467,7 +490,7 @@ impl OperationTranslation for SyncTranslation { _ => unreachable!(), }; require_input(&mut input, marker, args, results)?; - let mut e = crate::cute_sm100::Emit { + let mut e = crate::cute::Emit { session, location: input.location, ops: vec![], @@ -494,18 +517,18 @@ impl OperationTranslation for SyncTranslation { op.results = input.results; match self.0 { 0 => { - crate::cute_sm100::property(&mut op, "kind", "#nvvm")?; + crate::cute::property(&mut op, "kind", "#nvvm")?; } 1 => { - crate::cute_sm100::integer_property(&mut op, "size", 32, 4); - crate::cute_sm100::property(&mut op, "modifier", "#nvvm")?; + 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_sm100::property(&mut op, "kind", "#nvvm")?; + crate::cute::property(&mut op, "kind", "#nvvm")?; } _ => {} } 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 similarity index 95% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_sm120.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sm120.rs index d0fdb59947..8661fb7797 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_sm120.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sm120.rs @@ -44,6 +44,7 @@ pub fn register_nvvm_sm120_pack( Ok(()) } +/// Reorder reviewed register MMA operands and preserve packed lanes and scale selectors. struct MmaTranslation; impl OperationTranslation for MmaTranslation { fn translate( @@ -173,6 +174,7 @@ impl OperationTranslation for MmaTranslation { } } +/// Emit the exact matrix-load transpose, shape, packing, and result-register contract. struct LoadTranslation; impl OperationTranslation for LoadTranslation { fn translate( @@ -242,6 +244,7 @@ impl OperationTranslation for LoadTranslation { } } +/// Emit the verified matrix store without reordering its scalar register operands. struct StoreTranslation; impl OperationTranslation for StoreTranslation { fn translate( @@ -270,6 +273,7 @@ impl OperationTranslation for StoreTranslation { } } +/// Emit local shared-barrier TMA loads while rejecting hidden multicast or cache-policy changes. struct TmaTranslation; impl OperationTranslation for TmaTranslation { fn translate( @@ -353,6 +357,7 @@ impl OperationTranslation for TmaTranslation { } } +/// Check operand/result counts and canonical identity, then reject leftover source structure. pub(crate) fn require_input( input: &mut OperationInput, marker: &str, @@ -378,6 +383,7 @@ pub(crate) fn require_input( } Ok(()) } +/// Require an unsigned compile-time integer selector at the specified source operand. pub(crate) fn constant_integer( ctx: &Context, source: Ptr, @@ -399,6 +405,7 @@ pub(crate) fn constant_integer( .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, @@ -432,11 +439,13 @@ fn pointer( 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, @@ -458,6 +467,7 @@ fn extract_results( } Ok(()) } +/// Allocate one SSA identifier and return matching result and operand handles. pub(crate) fn fresh( session: &mut TranslationSession<'_>, ty: MlirType, @@ -465,6 +475,7 @@ pub(crate) fn fresh( 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, 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 similarity index 94% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_sregs.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sregs.rs index ec303a5849..e1c27f5a5c 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_sregs.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sregs.rs @@ -14,7 +14,10 @@ use dialect_nvvm::ops::{ ReadPtxSregNctaidYOp, ReadPtxSregNctaidZOp, ReadPtxSregNtidXOp, ReadPtxSregNtidYOp, ReadPtxSregNtidZOp, ReadPtxSregTidXOp, ReadPtxSregTidYOp, ReadPtxSregTidZOp, }; -use pliron::{context::Context, context::Ptr, operation::Operation}; +use pliron::{ + context::{Context, Ptr}, + operation::Operation, +}; use pliron_mlir_export::{ MlirOperation, OperationInput, OperationTranslation, TranslationError, TranslationRegistry, TranslationSession, @@ -73,8 +76,11 @@ pub fn register_nvvm_sreg_pack(registry: &mut TranslationRegistry) -> Result<(), 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, } @@ -122,6 +128,8 @@ impl OperationTranslation for SregTranslation { #[cfg(test)] mod tests { + //! One-to-one source/target checks for the reviewed coordinate-register subset. + use super::*; use crate::{ CutlassFullCuteMlir22, MlirConsumerProfile, @@ -141,6 +149,7 @@ mod tests { r#type::TypedHandle, }; + /// Append a chosen coordinate-read operation and return its source result. fn append_sreg( ctx: &mut Context, block: Ptr, @@ -163,6 +172,7 @@ mod tests { 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(); 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 similarity index 89% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_sync_restrict.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sync_restrict.rs index 8e43d67406..0cc70a98bf 100644 --- 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 @@ -21,6 +21,7 @@ use pliron_mlir_export::{ 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 { @@ -47,9 +48,13 @@ pub fn register_nvvm_sync_restrict_pack( 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>, } 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 similarity index 91% rename from cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_tcgen05.rs rename to cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/tcgen05.rs index 7db962bb7c..d193e790b9 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm_tcgen05.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/tcgen05.rs @@ -43,6 +43,7 @@ pub fn register_nvvm_tcgen05_pack( 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 { @@ -108,6 +109,7 @@ pub fn register_nvvm_tcgen05_pack( 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 { @@ -158,9 +160,13 @@ pub fn register_nvvm_tcgen05_pack( Ok(()) } +/// Preserve reviewed TCGen05 MMA operands, CTA group, shape, and A-collector selectors. struct MmaTranslation; -struct SharedOffsetTranslation(bool); +/// 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( @@ -287,25 +293,41 @@ impl OperationTranslation for MmaTranslation { } } +/// 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, } @@ -362,10 +384,15 @@ impl OperationTranslation for PrimitiveTranslation { /// 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, } @@ -487,6 +514,7 @@ impl OperationTranslation for MemoryTranslation { } } +/// Check canonical intrinsic identity, operand/result counts, and absence of leftover source structure. pub(crate) fn require_input( input: &mut OperationInput, marker: &str, @@ -513,10 +541,12 @@ pub(crate) fn require_input( 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, @@ -532,6 +562,7 @@ fn tmem_pointer( )) } +/// Convert the source address to CTA shared space before forming its native offset. fn shared_pointer( session: &mut TranslationSession<'_>, operations: &mut Vec, @@ -555,6 +586,7 @@ fn shared_pointer( Ok(converted) } +/// Allocate one SSA identifier and return matching result and operand handles. pub(crate) fn fresh( session: &mut TranslationSession<'_>, ty: MlirType, @@ -563,6 +595,7 @@ pub(crate) fn fresh( (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, @@ -578,6 +611,8 @@ pub(crate) fn operation( #[cfg(test)] mod tests { + //! Selector, identity, and register-order tests for tensor-memory instructions. + use super::*; use crate::{CutlassFullCuteMlir22, MlirConsumerProfile}; use dialect_mir::{ @@ -599,6 +634,7 @@ mod tests { }; use pliron_mlir_export::render_module; + /// Append a typed source operation to the fixture block and return its operation handle. fn append( ctx: &mut Context, block: Ptr, @@ -617,6 +653,7 @@ mod tests { 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); @@ -698,6 +735,7 @@ mod tests { (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 @@ -706,6 +744,7 @@ mod tests { Ok(render_module(&target)) } + /// 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(); @@ -746,6 +785,7 @@ mod tests { 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(); @@ -758,6 +798,7 @@ mod tests { ); } + /// 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] { @@ -784,12 +825,14 @@ mod tests { } } + /// 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( @@ -1054,6 +1097,7 @@ mod tests { } } + /// Reject a source instruction carrying a forged native catalog marker. #[test] fn intrinsic_identity_is_checked_before_export() { let (ctx, module, first) = fixture(); diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/packs.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/packs.rs index 30fbcba9c6..ce5807679c 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/packs.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/packs.rs @@ -3,6 +3,8 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Builtin scalar types, constants, and attributes shared by every export profile. + use pliron::{ attribute::AttrObj, builtin::{ @@ -54,6 +56,7 @@ pub fn register_builtin_pack(registry: &mut TranslationRegistry) -> Result<(), T Ok(()) } +/// Translate a builtin module and preserve its nested regions and optional symbol name. struct BuiltinModuleTranslation; impl OperationTranslation for BuiltinModuleTranslation { @@ -83,6 +86,7 @@ impl OperationTranslation for BuiltinModuleTranslation { // 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( @@ -126,6 +130,7 @@ impl OperationTranslation for BuiltinConstantTranslation { } } +/// Retain integer width while erasing source signedness from the MLIR type spelling. struct IntegerTypeTranslation; impl TypeTranslation for IntegerTypeTranslation { @@ -143,6 +148,7 @@ impl TypeTranslation for IntegerTypeTranslation { } } +/// Recursively translate function arguments and results through the active registry. struct FunctionTypeTranslation; impl TypeTranslation for FunctionTypeTranslation { @@ -169,6 +175,7 @@ impl TypeTranslation for FunctionTypeTranslation { } } +/// Represent Rust unit as an empty aggregate rather than inventing a scalar value. struct UnitTypeTranslation; impl TypeTranslation for UnitTypeTranslation { @@ -184,6 +191,7 @@ impl TypeTranslation for UnitTypeTranslation { } } +/// Translate a nested source type and preserve any registry diagnostic. fn translate_nested_type( ctx: &Context, registry: &TranslationRegistry, @@ -192,6 +200,7 @@ fn translate_nested_type( registry.translate_type(ctx, source) } +/// Preserve the bytes of a builtin string attribute. struct StringAttributeTranslation; impl AttributeTranslation for StringAttributeTranslation { @@ -208,6 +217,7 @@ impl AttributeTranslation for StringAttributeTranslation { } } +/// Represent a source identifier as the target string spelling. struct IdentifierAttributeTranslation; impl AttributeTranslation for IdentifierAttributeTranslation { @@ -224,6 +234,7 @@ impl AttributeTranslation for IdentifierAttributeTranslation { } } +/// Preserve a builtin boolean attribute's truth value. struct BoolAttributeTranslation; impl AttributeTranslation for BoolAttributeTranslation { @@ -241,6 +252,7 @@ impl AttributeTranslation for BoolAttributeTranslation { } } +/// Preserve integer width and all bits in the target signless attribute. struct IntegerAttributeTranslation; impl AttributeTranslation for IntegerAttributeTranslation { @@ -270,9 +282,13 @@ impl AttributeTranslation for IntegerAttributeTranslation { } } +/// 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, } @@ -316,6 +332,7 @@ impl AttributeTranslation for FloatAttributeTranslation { } } +/// Translate a source presence marker to an MLIR unit attribute. struct UnitAttributeTranslation; impl AttributeTranslation for UnitAttributeTranslation { @@ -329,6 +346,7 @@ impl AttributeTranslation for UnitAttributeTranslation { } } +/// Recursively translate the type carried by a source type attribute. struct TypeAttributeTranslation; impl AttributeTranslation for TypeAttributeTranslation { diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/primitives.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/primitives.rs index 90b1352d5e..1536f54587 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/primitives.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/primitives.rs @@ -29,12 +29,15 @@ 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") { @@ -64,12 +67,12 @@ impl MlirConsumerProfile for CutlassPrimitivesMlir22 { 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::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::nvvm::numeric::register_nvvm_profile_numeric_pack(&mut r, &self.gpu_arch)?; crate::register_nvvm_sm120_pack(&mut r)?; r.seal(); Ok(r) @@ -137,16 +140,22 @@ pub fn verify_primitives_source(ctx: &Context, module: &ModuleOp) -> Result<(), .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(()); @@ -193,6 +202,7 @@ impl SourceScan<'_> { } 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 { @@ -213,6 +223,7 @@ impl SourceScan<'_> { } Ok(()) } + /// Inspect location metadata for nested attributes carrying forbidden CuTe semantics. fn location(&mut self, loc: &Location) -> Result<(), ProfileError> { match loc { Location::Fused { diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/profile.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/profile.rs index 7442e906da..2d30d12ffb 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/src/profile.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/profile.rs @@ -3,6 +3,8 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Pinned consumer contracts, source validation, and executable GPU module packaging. + use pliron::{ builtin::ops::ModuleOp, context::Context, linked_list::ContainsLinkedList, op::Op, operation::Operation, @@ -24,45 +26,72 @@ use crate::{ /// 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, + /// 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(#[from] TranslationError), + 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(String), + 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(String), + 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(String), + 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(String), + 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()) } @@ -90,13 +119,17 @@ pub trait MlirConsumerProfile { /// 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 = "06d99a3df938394583715aacf5eadb65b4791c09"; + /// 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") { @@ -133,12 +166,12 @@ impl MlirConsumerProfile for CutlassFullCuteMlir22 { 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)?; + 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_sm100_attention::register_cute_sm100_attention_pack(&mut registry)?; + crate::cute::register_cute_sm100_attention_pack(&mut registry)?; registry.seal(); Ok(registry) } @@ -218,7 +251,7 @@ pub(crate) fn add_cutlass_executable_envelope(module: &mut MlirModule) -> Result 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)?; + crate::mir::shared::hoist_shared_globals(&mut device_operations)?; let mut gpu_module = MlirOperation::new("gpu.module")?; gpu_module @@ -244,7 +277,9 @@ pub(crate) fn add_cutlass_executable_envelope(module: &mut MlirModule) -> Result 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 { @@ -287,9 +322,12 @@ pub(crate) fn render_mapping_module_without_cutlass_envelope( #[cfg(test)] mod tests { + //! Profile manifest, architecture rejection, and immutable-source verification checks. + use super::*; use pliron::identifier::Identifier; + /// 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(); @@ -300,6 +338,7 @@ mod tests { 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!( @@ -308,6 +347,7 @@ mod tests { )); } + /// 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(); diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/stack-temporaries.md b/cuda-oxide/crates/cuda-oxide-mlir-export/stack-temporaries.md deleted file mode 100644 index 59c8d01306..0000000000 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/stack-temporaries.md +++ /dev/null @@ -1,35 +0,0 @@ -# Why move stack temporaries to the function entry? - -LLVM can often replace temporary memory with ordinary values. Its scalar -replacement (SROA) and memory-to-register (mem2reg) passes look for allocations -in the function's entry block. - -Simple pseudocode: - -```text -Before: - for each x: - temp = stack_alloc() - store(temp, x + 1) - consume(load(temp)) - -After hoisting: - temp = stack_alloc() - for each x: - store(temp, x + 1) - consume(load(temp)) - -After LLVM optimizes: - for each x: - consume(x + 1) -``` - -Only the allocation moves initially. Computation, writes and reads stay in the -loop. LLVM can then remove the unnecessary store/load pair and the unused slot. -Entry placement enables this optimization; it does not guarantee elimination. - -Our `hoist_enum_scratch` in [`mir_checked.rs`](src/mir_checked.rs) applies this -to fixed-size, compiler-owned enum temporaries whose addresses never escape. -It preserves checks and enum layout, and leaves ordinary Rust allocations alone. - -See [LLVM's guidance on allocas](https://llvm.org/docs/Frontend/PerformanceTips.html#use-of-allocas). 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 index 148782fbb4..43405d5ea5 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/address_contract.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/address_contract.rs @@ -1,3 +1,5 @@ +//! 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}; @@ -22,6 +24,7 @@ use pliron::{ }; use pliron_mlir_export::{MlirOperation, MlirType}; +/// Build pointer-address operations for a selected source space and export through the primitives profile. fn translate( space: u32, width: u32, @@ -81,6 +84,7 @@ fn translate( .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] { @@ -127,6 +131,7 @@ fn explicit_shared_offsets_preserve_cluster_rank_while_rust_addresses_genericize } } +/// 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] { 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 index bc04f409df..499569dd39 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_cute.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_cute.rs @@ -2,6 +2,8 @@ * 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::{ @@ -22,6 +24,7 @@ use pliron::{ 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] { 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 index 28fd4088ef..5a4313fd79 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_descriptor.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_descriptor.rs @@ -27,12 +27,17 @@ use pliron::{ use pliron_mlir_export::{MlirAttribute, MlirBlock, MlirOperation, MlirType, MlirValueId}; 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 { @@ -54,6 +59,7 @@ fn function_block(op: &MlirOperation) -> Option<&MlirBlock> { } 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); @@ -128,6 +134,7 @@ impl Probe { } } + /// 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; @@ -154,6 +161,7 @@ impl Probe { } } +/// 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); @@ -173,6 +181,7 @@ fn descriptor_update_preserves_metadata_across_the_entire_address_field() { } } +/// 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) { 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 index fc11c3e4a0..d4d5a2ebb8 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_primitives.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_primitives.rs @@ -20,6 +20,7 @@ use pliron::{ 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(), @@ -53,6 +54,7 @@ fn mark(ctx: &mut Context, op: Ptr, block: Ptr, marker: & } } +/// 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(); 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 index d2098c79df..89e9fcf3d7 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/checked_storage.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/checked_storage.rs @@ -1,3 +1,5 @@ +//! 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::{CutlassPrimitivesMlir22, MlirConsumerProfile}; @@ -22,6 +24,7 @@ use pliron::{ value::Value, }; use pliron_mlir_export::render_module; +/// Append a typed source operation to the fixture block and return its operation handle. fn append( ctx: &mut Context, b: Ptr, @@ -32,6 +35,7 @@ fn append( 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, @@ -54,6 +58,7 @@ fn function( 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() @@ -61,6 +66,7 @@ fn render(ctx: &Context, m: &ModuleOp) -> Result { .map(|m| render_module(&m)) .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] { @@ -107,6 +113,7 @@ fn checked_arithmetic_preserves_signedness_and_padded_overflow_tuple() { } } } +/// 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(); @@ -151,6 +158,7 @@ fn direct_enum_uses_physical_offsets_for_tag_and_active_payload() { ); 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); @@ -198,6 +206,7 @@ fn shared_case(conflict: &str) -> Result { append::(&mut ctx, b, vec![], vec![]); render(&ctx, &m) } +/// Check matching shared declarations merge, conflicting ones fail, and zero alignment uses the natural value. #[test] fn shared_declarations_deduplicate_only_identical_storage_and_zero_means_natural() { let text = shared_case("same").unwrap(); @@ -221,6 +230,7 @@ fn shared_declarations_deduplicate_only_identical_storage_and_zero_means_natural } } +/// 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(); @@ -258,6 +268,7 @@ fn primitive_dynamic_shared_uses_external_as3_storage_with_alignment_and_offsets ); } +/// 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(); @@ -279,6 +290,7 @@ fn primitive_contract_failure_emits_trap_before_unreachable() { ); } +/// 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(); @@ -307,6 +319,7 @@ fn logical_unit_call_uses_void_signature_and_preserves_value_uses() { 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(); @@ -333,6 +346,7 @@ fn explicitly_returned_zero_sized_aggregates_keep_their_call_result() { } } +/// 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}; @@ -359,6 +373,7 @@ fn rust_scalar_intrinsics_are_operations_not_external_calls() { } } +/// 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(); @@ -383,6 +398,7 @@ fn unrolling_builtin_integer_constants_retain_width_and_bits() { 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(); @@ -412,6 +428,7 @@ fn rust_fp32_max_preserves_nan_ignoring_and_subnormal_contract() { ); } +/// 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)] { 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 index b05ae2cebb..28de95c500 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs @@ -3,6 +3,8 @@ * 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::{CutlassFullCuteMlir22, MlirConsumerProfile}; @@ -35,12 +37,14 @@ use pliron::{ }; use pliron_mlir_export::render_module; +/// 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(); @@ -262,6 +266,7 @@ fn all_eight_elementwise_ops_map_to_cutlass_cute_and_core_mlir() { 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(); 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 index 9e42f6c4bb..23c1b7a884 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/cute_gemv.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/cute_gemv.rs @@ -3,6 +3,8 @@ * 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 dialect_cute::{ attributes::CuteTensorRoleAttr, @@ -31,12 +33,14 @@ use pliron::{ }; use pliron_mlir_export::{MlirModule, render_module}; +/// 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); @@ -216,6 +220,7 @@ fn build_gemv_module(ctx: &mut Context) -> ModuleOp { 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() @@ -223,10 +228,12 @@ fn translate_gemv_target(ctx: &Context, module: &ModuleOp) -> MlirModule { .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)) } +/// 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(); 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 index 03a5f851a9..7a9c7ab395 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/grid_constant.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/grid_constant.rs @@ -3,6 +3,8 @@ * 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}; @@ -30,6 +32,7 @@ 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(); @@ -145,6 +148,7 @@ fn fixture(case: &str) -> Result<(String, Vec), String> { Ok((render_module(&target), 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(); @@ -170,6 +174,7 @@ fn descriptors_remain_by_value_and_grid_constant_after_slice_flattening() { ); } +/// 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 [ 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 index 07eefd33b0..c7f1485f5c 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/nvvm_attention.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/nvvm_attention.rs @@ -1,3 +1,5 @@ +//! 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 @@ -20,13 +22,18 @@ use pliron::{ }; use pliron_mlir_export::render_module; +/// 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 { let mut ctx = Context::new(); dialect_mir::register(&mut ctx); @@ -52,10 +59,12 @@ impl Fixture { 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( @@ -67,6 +76,7 @@ impl Fixture { 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( @@ -85,6 +95,7 @@ impl Fixture { 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, @@ -103,6 +114,7 @@ impl Fixture { self.append(op, marker) } + /// Append the source function return after the tested instructions. fn finish(&mut self) { let op = Operation::new( &mut self.ctx, @@ -115,10 +127,12 @@ impl Fixture { 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) @@ -127,6 +141,7 @@ impl Fixture { } } +/// 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] = [ @@ -156,6 +171,7 @@ fn ftz_exp2_keeps_its_exact_subnormal_mode_in_both_profiles() { } } +/// 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] = [ @@ -180,6 +196,7 @@ fn ftz_exp2_rejects_other_canonical_identities_in_both_profiles() { } } +/// 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] = [ @@ -221,6 +238,7 @@ fn ftz_exp2_rejects_wrong_type_and_selectors() { } } +/// 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(); @@ -251,6 +269,7 @@ fn softmax_math_keeps_approximation_and_subnormal_modes() { 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(); @@ -273,6 +292,7 @@ fn unsupported_math_and_forged_identity_are_rejected() { ); } +/// 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(); @@ -296,6 +316,7 @@ fn butterfly_shuffle_preserves_mask_offset_and_warp_clamp() { ); } +/// 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(); @@ -324,6 +345,7 @@ fn cp_async_commit_and_wait_keep_the_group_count() { ); } +/// 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)] { @@ -339,6 +361,7 @@ fn cp_async_wait_rejects_runtime_or_out_of_range_counts() { } } +/// 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(); 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 index f8818aac46..65d286e510 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/primitives_profile.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/primitives_profile.rs @@ -3,12 +3,15 @@ * 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}; @@ -33,6 +36,7 @@ fn primitives_rejects_cute_hidden_in_pointer_type_metadata() { 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(); @@ -50,6 +54,7 @@ fn primitives_profile_translates_without_cute_registration() { 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(); @@ -73,6 +78,7 @@ fn primitives_profile_pins_consumer_and_target() { } } +/// 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(); @@ -95,6 +101,7 @@ fn primitives_rejects_cute_operations_before_local_verification() { 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}; @@ -119,6 +126,7 @@ fn primitives_rejects_cute_attributes_in_dictionaries() { 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(); @@ -141,6 +149,7 @@ fn primitives_rejects_unmapped_nvvm_instructions() { 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(); 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 index 5577f783be..4ce520802a 100644 --- 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 @@ -1,3 +1,5 @@ +//! 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 @@ -21,6 +23,7 @@ use pliron::{ }; use pliron_mlir_export::render_module; +/// Build and export a scalar call with independently controlled operand and result signatures. fn translate_call( profile: &dyn MlirConsumerProfile, callee: &str, @@ -87,6 +90,7 @@ fn translate_call( .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()), @@ -94,6 +98,7 @@ fn profiles() -> [Box; 2] { ] } +/// 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() { @@ -118,6 +123,7 @@ fn rust_f32_max_exports_nan_ignoring_subnormal_preserving_native_maximum() { } } +/// 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() { @@ -138,6 +144,7 @@ fn rust_f32_max_rejects_malformed_operand_and_result_counts() { } } +/// 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() { @@ -159,6 +166,7 @@ fn rust_f32_max_rejects_non_f32_operands_and_results() { } } +/// 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() { 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 index 2b82d673f7..33da141882 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/sm120_primitives.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/sm120_primitives.rs @@ -1,3 +1,5 @@ +//! 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 @@ -24,13 +26,19 @@ use pliron::{ }; use pliron_mlir_export::render_module; +/// 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); @@ -85,9 +93,11 @@ impl Fixture { 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, @@ -111,6 +121,7 @@ impl Fixture { 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()]); @@ -123,9 +134,11 @@ impl Fixture { ); 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() @@ -133,6 +146,7 @@ impl Fixture { .map(|m| render_module(&m)) .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( @@ -152,6 +166,7 @@ impl Fixture { 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. @@ -176,6 +191,7 @@ impl Fixture { } } +/// 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(); @@ -203,6 +219,7 @@ fn sm120_mxf4_preserves_a_b_c_scale_order_and_extracts_all_results() { 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(); @@ -226,6 +243,7 @@ fn sm120_mxf4_rejects_other_variants_and_forged_identity() { ); } +/// Reject the SM120-only MXF4 instruction when targeting SM100. #[test] fn sm100_rejects_sm120_mxf4() { let mut f = Fixture::new(); @@ -251,6 +269,7 @@ fn sm100_rejects_sm120_mxf4() { ); } +/// 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(); @@ -329,6 +348,7 @@ fn fp16_mma_preserves_packed_a_b_and_f32_accumulator_order() { ); } +/// 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 [ @@ -390,6 +410,7 @@ fn matrix_loads_preserve_transposition_and_register_count() { } } +/// 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(); @@ -410,6 +431,7 @@ fn sm120_load_pack_and_store_preserve_register_and_lane_order() { 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(); @@ -447,6 +469,7 @@ fn sm120_rejects_other_matrix_load_selectors_and_forged_identity() { ); } +/// 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(); @@ -471,6 +494,7 @@ fn sm120_tma_uses_local_shared_barrier_and_no_multicast() { 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)] { @@ -491,6 +515,7 @@ fn sm120_tma_rejects_nonzero_multicast_mask_or_cache_hint() { } } +/// Reject SM100-only tensor-memory instructions during source architecture validation. #[test] fn sm120_rejects_tcgen05_before_translation() { let mut f = Fixture::new(); @@ -512,6 +537,7 @@ fn sm120_rejects_tcgen05_before_translation() { ); } +/// 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(); 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 index 7b1867b257..205682759b 100644 --- a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/sync_restrict.rs +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/sync_restrict.rs @@ -1,3 +1,5 @@ +//! 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 @@ -21,14 +23,20 @@ use pliron::{ }; use pliron_mlir_export::render_module; +/// 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); @@ -75,6 +83,7 @@ impl Fixture { } } + /// 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") @@ -89,6 +98,7 @@ impl Fixture { } } +/// 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(); @@ -116,6 +126,7 @@ fn selected_cutlass_profiles_preserve_proxy_direction_and_each_memory_order() { } } +/// Reject fences carrying the other direction's catalog identity. #[test] fn proxy_fences_reject_swapped_intrinsic_markers() { for acquire in [false, true] { @@ -137,6 +148,7 @@ fn proxy_fences_reject_swapped_intrinsic_markers() { } } +/// Reject source attributes the reviewed fence recipe has not explicitly consumed. #[test] fn proxy_fences_reject_unexpected_semantic_attributes() { let fixture = Fixture::new(); @@ -148,6 +160,7 @@ fn proxy_fences_reject_unexpected_semantic_attributes() { 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() { 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 similarity index 93% rename from cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled_emit.rs rename to cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled/emit.rs index e3c222f231..6119465339 100644 --- 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 @@ -34,8 +34,9 @@ use pliron::r#type::{TypeHandle, Typed}; use pliron::value::Value; use rustc_public::mir; -use super::block_scaled::{BlockScaledStage, decode_block_scaled}; +use super::super::block_scaled::{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, @@ -49,6 +50,7 @@ fn insert( } } +/// Branch from the last emitted effect to the call continuation; reject diverging calls. fn finish( ctx: &mut Context, effect: Ptr, @@ -66,6 +68,7 @@ fn finish( 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, @@ -93,6 +96,7 @@ fn finish_value( 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, @@ -116,6 +120,7 @@ fn extract_field( (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, @@ -150,6 +155,7 @@ fn split_u8_slice( Ok((data, len, len_op)) } +/// Require the destination's scaled-view stage and return its matrix/vector role. fn destination_facts( body: &mir::Body, destination: &mir::Place, @@ -182,6 +188,7 @@ fn destination_facts( Ok(decoded.role) } +/// 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, @@ -268,6 +275,7 @@ pub(crate) fn emit_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, @@ -321,6 +329,7 @@ pub(crate) fn emit_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, @@ -371,6 +380,7 @@ pub(crate) fn emit_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, @@ -420,6 +430,7 @@ pub(crate) fn emit_load( ) } +/// 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, diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled/mod.rs similarity index 79% rename from cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled.rs rename to cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled/mod.rs index 74734b7c5c..7e49f264c1 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled/mod.rs @@ -17,6 +17,8 @@ //! 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, @@ -31,32 +33,48 @@ use rustc_public::ty::{GenericArgs, RigidTy, Ty, TyKind}; use super::tensor::{canonical_adt_path, const_arg, exact_schema, type_arg, unsigned_const}; -const BLOCK_SCALED_TENSOR: &str = "cute_rs::block_scaled::BlockScaledTensor"; -const BLOCK_SCALED_THREAD_ROW: &str = "cute_rs::block_scaled::BlockScaledThreadRow"; -const BLOCK_SCALED_TILE_64: &str = "cute_rs::block_scaled::BlockScaledTile64"; -const LOADED_BLOCK_SCALED_TILE_64: &str = "cute_rs::block_scaled::LoadedBlockScaledTile64"; -const K_MAJOR: &str = "cute_rs::block_scaled::KMajor"; +/// 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::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::E2M1`; public reexports do not change it. const E2M1: &str = "cute_rs::numeric::E2M1"; +/// Exact defining path of `cute_rs::numeric::UE8M0`; public reexports do not change it. const UE8M0: &str = "cute_rs::numeric::UE8M0"; /// 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, + /// 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; @@ -64,6 +82,7 @@ fn adt_path(ty: &Ty) -> Option<(String, GenericArgs)> { 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")); @@ -74,6 +93,7 @@ fn exact_marker(ty: &Ty, expected: &str) -> Result<(), String> { 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()); @@ -88,6 +108,7 @@ fn decode_role(ty: &Ty) -> Result { } } +/// 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()); @@ -161,6 +182,7 @@ pub(crate) fn decode_block_scaled_receiver(ty: &Ty) -> Result(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, "tiled_copy") == 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, "block_scaled_mma::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, "block_scaled_mma::fragments").filter(|method| { + matches!(*method, "zero" | "pairs_at_unchecked" | "accumulate_k64") + }) + })?; + 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) + } + "accumulate_k64" + 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") { + 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, "block_scaled_mma::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), + "slot" 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) + } + "current_tile" + if args.len() == 1 + && matches!((receiver_scheduler, destination_tile), (Some(source), Some(result)) if source == result) => + { + Some(CuteFn::SchedulerCurrentTile) + } + "coordinates" if args.len() == 1 && receiver_tile.is_some() => { + Some(CuteFn::WorkTileCoordinates) + } + "advance" 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, + }, + "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) + } + "store_linear" + 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", + "block_scaled_mma::tiled", + "block_scaled_mma::fragments", + "block_scaled_mma::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..3be5f7a744 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/dispatch.rs @@ -0,0 +1,751 @@ +/* + * 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 be cute calls (indirect dispatch of cute_rs items is rejected by + // `requires_direct_dispatch` in the function-item path). + 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::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( + "TmaLoadPipeline::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 index 297c905266..90f16af760 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/emit.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/emit.rs @@ -33,7 +33,7 @@ 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::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, @@ -551,6 +551,7 @@ pub(crate) fn emit_load_matrix( 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. @@ -562,6 +563,7 @@ mod tests { 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, @@ -575,6 +577,7 @@ mod tests { (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) @@ -584,6 +587,7 @@ mod tests { .address_space } + /// Insert an undefined carrier after the test's current operation. fn undef_after( ctx: &mut Context, block: Ptr, @@ -595,6 +599,7 @@ mod tests { (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}; @@ -638,6 +643,7 @@ mod tests { 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(); @@ -692,6 +698,7 @@ mod tests { 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(); @@ -1454,6 +1461,7 @@ pub(crate) fn emit_store_global_no_allocate( 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, 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 similarity index 92% rename from cuda-oxide/crates/mir-importer/src/translator/cute/epilogue_emit.rs rename to cuda-oxide/crates/mir-importer/src/translator/cute/epilogue/emit.rs index 45f6257094..8aa4ab49a0 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/epilogue_emit.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/epilogue/emit.rs @@ -31,22 +31,24 @@ use pliron::r#type::TypeHandle; use pliron::value::Value; use rustc_public::mir; -use super::emit::recover_cta_shared_pointer; -use super::epilogue::{ +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::smem_mma_emit::{ +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(()), @@ -57,6 +59,7 @@ fn require_type_fact(result: Result, what: &str) -> TranslationRes } } +/// Find an aggregate field's translated carrier type before extracting it. fn translated_field_type( ctx: &mut Context, body: &mir::Body, @@ -69,6 +72,7 @@ fn translated_field_type( 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, @@ -90,6 +94,7 @@ fn epilogue_base( 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, @@ -115,6 +120,7 @@ fn warp_slice_carriers( 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, @@ -184,6 +190,7 @@ pub(crate) fn emit_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, @@ -266,6 +273,7 @@ pub(crate) fn emit_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, @@ -318,6 +326,7 @@ pub(crate) fn emit_store_fragment( 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, @@ -356,6 +365,7 @@ pub(crate) fn emit_sync( 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, @@ -424,13 +434,18 @@ pub(crate) fn emit_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, diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/epilogue.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/epilogue/mod.rs similarity index 79% rename from cuda-oxide/crates/mir-importer/src/translator/cute/epilogue.rs rename to cuda-oxide/crates/mir-importer/src/translator/cute/epilogue/mod.rs index 75d9060764..4de60fcce8 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/epilogue.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/epilogue/mod.rs @@ -9,6 +9,8 @@ //! 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; @@ -20,11 +22,16 @@ 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::Sm120Epilogue128x128`; public reexports do not change it. const EPILOGUE_TILE: &str = "cute_rs::epilogue::Sm120Epilogue128x128"; +/// Exact defining path of `cute_rs::epilogue::Sm120EpilogueWarp128x128`; public reexports do not change it. const EPILOGUE_WARP_SLICE: &str = "cute_rs::epilogue::Sm120EpilogueWarp128x128"; -const ACCUMULATOR: &str = "cute_rs::block_scaled_mma::Mxf4AccumulatorTile2x8"; +/// Exact defining path of `cute_rs::block_scaled_mma::fragments::Mxf4AccumulatorTile2x8`; public reexports do not change it. +const ACCUMULATOR: &str = "cute_rs::block_scaled_mma::fragments::Mxf4AccumulatorTile2x8"; +/// Exact defining path of `cute_rs::pipeline::TmaStorePipeline`; public reexports do not change it. const TMA_STORE_PIPELINE: &str = "cute_rs::pipeline::TmaStorePipeline"; +/// 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); @@ -35,6 +42,7 @@ fn exact_adt(ty: &Ty, expected: &str) -> Result, String> { 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, @@ -47,6 +55,7 @@ fn behind_one_pointer( } } +/// 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); @@ -55,26 +64,32 @@ fn is_exact_marker(ty: &Ty, expected: &str) -> Result { 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 { @@ -127,6 +142,7 @@ pub(crate) fn one_unsigned_const(func: &mir::Operand, what: &str) -> Result TypeHandle { let storage: TypeHandle = MirFP16Type::get(ctx).into(); let mma = CuteTiledMmaPlanAttr::mxf4_128x128x128(canonical_data_placement()); @@ -134,10 +150,12 @@ pub(crate) fn tile_type(ctx: &mut Context) -> TypeHandle { 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(); 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..c6ab1edeed --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/identity.rs @@ -0,0 +1,374 @@ +/* + * 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", + "block_scaled", + "block_scaled_mma", + "epilogue", + "numeric", + "pipeline", + "scheduler", + "sm100", + "sm100_attention", + "tensor", + "tiled_copy", +]; + +/// 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; + }; + let first_module = rest.split("::").next().unwrap_or_default(); + CUTE_SOURCE_ABSTRACTION_MODULES.contains(&first_module) +} + +/// Rebuild the canonical definition path by walking DefId parents, +/// mirroring `intrinsics::generated::classify_raw_intrinsic`. +pub(super) fn canonical_path(fn_def: &FnDef, crate_name: &str) -> String { + let mut segments = Vec::new(); + let mut current = Some(fn_def.def_id()); + while let Some(def_id) = current { + let printed = def_id.name(); + let segment = printed.as_str().rsplit("::").next().unwrap_or_default(); + segments.push(segment.to_owned()); + current = def_id.parent(); + } + canonical_path_from_leaf_segments(crate_name, segments) +} + +/// Join a DefId parent walk without erasing a same-named nested module. +pub(in crate::translator) fn canonical_path_from_leaf_segments( + crate_name: &str, + mut segments: Vec, +) -> String { + segments.reverse(); + // Only the first segment is the crate root. A nested module may legally + // have the same spelling and must remain part of the identity: + // + // cute_rs::cute_rs::tile::load_tile != cute_rs::tile::load_tile + if segments + .first() + .is_some_and(|segment| segment == crate_name) + { + segments.remove(0); + } + format!("{crate_name}::{}", segments.join("::")) +} + +/// 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::{canonical_path_from_leaf_segments, is_source_abstraction_path}; + + /// Check that path normalization removes only the crate-root segment. + #[test] + fn removes_only_the_crate_root_segment() { + let path = canonical_path_from_leaf_segments( + "cute_rs", + ["load_tile", "tile", "cute_rs", "cute_rs"] + .into_iter() + .map(str::to_string) + .collect(), + ); + assert_eq!(path, "cute_rs::cute_rs::tile::load_tile"); + } + + /// 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::tile::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::assume_warp_uniform", + "cute_rs" + )); + assert_eq!( + super::CUTE_SOURCE_ABSTRACTION_MODULES, + &[ + "arch", + "block_scaled", + "block_scaled_mma", + "epilogue", + "numeric", + "pipeline", + "scheduler", + "sm100", + "sm100_attention", + "tensor", + "tiled_copy", + ] + ); + assert!(is_source_abstraction_path( + "cute_rs::block_scaled::tensor::BlockScaledTensor<...>::thread_row", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute-rs::numeric::PackedE2M1x2::to_f32x2", + "cute-rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::tensor::Tensor<...>::layout", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::tiled_copy::TiledCopy<...>::copy", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::block_scaled_mma::tiled::Mxfp4TiledMma<...>::gemm", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::epilogue::Sm120EpilogueWarp128x128::store_atom", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::pipeline::TmaLoadPipeline<...>::consumer_wait", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::scheduler::StaticPersistentTileScheduler<...>::current", + "cute_rs" + )); + assert!(!is_source_abstraction_path( + "cute_rs::tile::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" + )); + } + + /// 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_linear"), + Some("store_linear") + ); + 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}::coordinates"), + Some("coordinates") + ); + 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::TmaLoadPipeline<...>>::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::block_scaled_mma::tiled::Mxfp4TiledMma<...>>::load_a_128", + "block_scaled_mma::tiled" + ), + Some("load_a_128") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute-rs::tiled_copy::{impl#7}::from_raw_parts", + "tiled_copy" + ), + Some("from_raw_parts") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::other::Mxfp4TiledMma<...>>::load_a_128", + "block_scaled_mma::tiled" + ), + None + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::block_scaled_mma::load_a_128", + "block_scaled_mma::tiled" + ), + None, + "a free function is not an inherent method" + ); + } + + /// 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::Sm120Epilogue128x128::sync_ready_for_tma", + "epilogue" + ), + Some("sync_ready_for_tma") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::block_scaled_mma::epilogue::{impl#9}::store_tile", + "block_scaled_mma::epilogue" + ), + Some("store_tile") + ); + assert_eq!( + pipeline_method_name_from_path( + "cute_rs::pipeline::TmaStorePipeline<1>::producer_commit" + ), + Some("producer_commit") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::epilogue::Sm120EpilogueWarp128x128::store_tile", + "block_scaled_mma::epilogue" + ), + None, + "store_tile is defined in block_scaled_mma::epilogue, not epilogue" + ); + } +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/mod.rs index cd3fad21ca..a432ad2433 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/mod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/mod.rs @@ -17,2283 +17,36 @@ //! definition path: cute_rs::tile::load_tile -> CUTE_FNS entry //! ``` //! -//! A recognized call becomes a CuTe operation; its Rust body is not translated. -//! The importer checks its specialized generic arguments against [`SubstKind`] -//! before reading static configuration. Ordinary Rust helpers in explicitly -//! allowed modules use normal body translation and are inlined by LLVM. -//! Any other unknown `cute_rs` call is an error: translating a compiler stub's -//! body could turn layout operations into scalar code and lose their meaning. +//! 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 block_scaled_emit; pub(crate) mod emit; pub(crate) mod epilogue; -pub(crate) mod epilogue_emit; pub(crate) mod layout; pub(crate) mod pipeline; -pub(crate) mod pipeline_emit; pub(crate) mod scheduler; -pub(crate) mod scheduler_emit; -pub(crate) mod sm100_attention_emit; -pub(crate) mod sm100_emit; +pub(crate) mod sm100; +pub(crate) mod sm100_attention; pub(crate) mod smem_mma; -pub(crate) mod smem_mma_emit; pub(crate) mod static_config; pub(crate) mod tensor; -pub(crate) mod tensor_emit; - -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 rustc_public::ty::FnDef; - -/// Recognized cute-rs entry points. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum CuteFn { - Sm100(sm100_emit::Sm100Fn), - Attention(sm100_attention_emit::AttentionFn), - MakeTensorRead, - MakeTensorWrite, - ZippedDivideRead, - ZippedDivideWrite, - SliceRead, - SliceWrite, - TensorTileIsFull, - TensorTileBase, - TensorLoadTile, - TensorStoreTile, - TensorStoreElementAbs, - BlockScaledMake, - BlockScaledThreadRow, - BlockScaledKTile, - BlockScaledLoadK64, - BlockScaledDotK64, - SchedulerNew1d, - SchedulerHasWork, - SchedulerCurrentTile, - WorkTileCoordinates, - SchedulerAdvance, - TmaLoadPipelineMake, - TmaLoadPipelineFromRawPartsRejected, - TmaLoadPipelineInit, - PipelineStateNew, - PipelineStateSlot, - PipelineStateAdvance, - PipelineProducerAcquire, - PipelineProducerExpectTx, - PipelineConsumerWait, - PipelineConsumerRelease, - PipelineProducerTail, - SharedTensorOverlay, - TiledMmaSlice, - MmaFragmentFill, - MmaLoadScales, - MmaFragmentSliceK, - MmaLoadA, - MmaPartitionB, - TiledGemm, - EpilogueOverlay, - EpilogueWarpSlice, - EpilogueStoreFragment, - EpilogueSyncReusable, - EpilogueSyncReadyForTma, - EpilogueHalf, - TmaStoreAcquire, - TmaStoreCommit, - TmaStoreTail, - LoadTile, - StoreTile, - StoreGlobalNoAllocate, - AssumeDiv, - CopyG2S, - LoadMatrixA, - LoadMatrixB, - LoadMatrixBPair, - LoadMatrixBPairKMajor, - CopyTma2d, - 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 { - Type, - Const, - Lifetime, -} - -/// A recognized function's definition path and expected generic arguments. -/// Re-export paths do not appear here. -pub(crate) struct CuteFnSpec { - pub canonical: &'static str, - pub func: CuteFn, - pub schema: &'static [SubstKind], -} - -/// One entry per recognized function, including its generic argument schema. -pub(crate) const CUTE_FNS: &[CuteFnSpec] = &[ - CuteFnSpec { - canonical: "cute_rs::sm100_attention::__compiler::mma_ss", - func: CuteFn::Attention(sm100_attention_emit::AttentionFn::MmaSs), - schema: &[ - SubstKind::Type, - SubstKind::Type, - SubstKind::Const, - SubstKind::Const, - SubstKind::Const, - SubstKind::Const, - ], - }, - CuteFnSpec { - canonical: "cute_rs::sm100_attention::__compiler::mma_ts", - func: CuteFn::Attention(sm100_attention_emit::AttentionFn::MmaTs), - schema: &[ - SubstKind::Type, - SubstKind::Const, - SubstKind::Const, - SubstKind::Const, - SubstKind::Const, - ], - }, - CuteFnSpec { - canonical: "cute_rs::sm100_attention::__compiler::tma_load", - func: CuteFn::Attention(sm100_attention_emit::AttentionFn::TmaLoad), - schema: &[SubstKind::Type, SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100_attention::__compiler::tma_store", - func: CuteFn::Attention(sm100_attention_emit::AttentionFn::TmaStore), - schema: &[SubstKind::Type, SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_tma_store", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::TmaStore), - schema: &[SubstKind::Type], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_store_commit", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::StoreCommit), - schema: &[SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_store_acquire", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::StoreAcquire), - schema: &[SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_store_tail", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::StoreTail), - schema: &[SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_tmem_alloc", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::TmemAlloc), - schema: &[SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_tmem_dealloc", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::TmemDealloc), - schema: &[SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_tiled_mma", - func: CuteFn::Sm100(sm100_emit::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_emit::Sm100Fn::TmemEpilogue), - schema: &[SubstKind::Type], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_cluster_tma_load", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::ClusterTmaLoad), - schema: &[SubstKind::Type], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_pipeline_init", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::PipelineInit), - schema: &[SubstKind::Const, SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_pipeline_acquire", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::PipelineAcquire), - schema: &[SubstKind::Const, SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_pipeline_expect", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::PipelineExpect), - schema: &[SubstKind::Const, SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_pipeline_wait", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::PipelineWait), - schema: &[SubstKind::Const, SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_pipeline_release", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::PipelineRelease), - schema: &[SubstKind::Const, SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_accumulator_init", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::AccumulatorInit), - schema: &[SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_accumulator_acquire", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::AccumulatorAcquire), - schema: &[SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_accumulator_commit", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::AccumulatorCommit), - schema: &[SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_accumulator_wait", - func: CuteFn::Sm100(sm100_emit::Sm100Fn::AccumulatorWait), - schema: &[SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::sm100::__compiler::sm100_accumulator_release", - func: CuteFn::Sm100(sm100_emit::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::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::tiled_copy::__compiler::shared_tensor_overlay", - func: CuteFn::SharedTensorOverlay, - schema: &[ - SubstKind::Type, - SubstKind::Type, - SubstKind::Type, - SubstKind::Type, - ], - }, - CuteFnSpec { - canonical: "cute_rs::block_scaled_mma::__compiler::tiled_mma_slice", - func: CuteFn::TiledMmaSlice, - schema: &[SubstKind::Type], - }, - CuteFnSpec { - canonical: "cute_rs::block_scaled_mma::__compiler::fragment_fill", - func: CuteFn::MmaFragmentFill, - schema: &[], - }, - CuteFnSpec { - canonical: "cute_rs::block_scaled_mma::__compiler::mma_load_scales", - func: CuteFn::MmaLoadScales, - schema: &[SubstKind::Type], - }, - CuteFnSpec { - canonical: "cute_rs::block_scaled_mma::__compiler::fragment_slice_k", - func: CuteFn::MmaFragmentSliceK, - schema: &[], - }, - CuteFnSpec { - canonical: "cute_rs::block_scaled_mma::__compiler::mma_load_a", - func: CuteFn::MmaLoadA, - schema: &[SubstKind::Type], - }, - CuteFnSpec { - canonical: "cute_rs::block_scaled_mma::__compiler::mma_partition_b", - func: CuteFn::MmaPartitionB, - schema: &[SubstKind::Type], - }, - CuteFnSpec { - canonical: "cute_rs::block_scaled_mma::__compiler::tiled_gemm", - func: CuteFn::TiledGemm, - schema: &[SubstKind::Type], - }, - CuteFnSpec { - canonical: "cute_rs::epilogue::__compiler::epilogue_smem_overlay", - func: CuteFn::EpilogueOverlay, - schema: &[], - }, - CuteFnSpec { - canonical: "cute_rs::epilogue::__compiler::epilogue_warp_slice", - func: CuteFn::EpilogueWarpSlice, - schema: &[], - }, - CuteFnSpec { - canonical: "cute_rs::block_scaled_mma::__compiler::epilogue_store_fragment", - func: CuteFn::EpilogueStoreFragment, - schema: &[], - }, - CuteFnSpec { - canonical: "cute_rs::epilogue::__compiler::epilogue_sync_reusable", - func: CuteFn::EpilogueSyncReusable, - schema: &[], - }, - CuteFnSpec { - canonical: "cute_rs::epilogue::__compiler::epilogue_sync_ready_for_tma", - func: CuteFn::EpilogueSyncReadyForTma, - schema: &[], - }, - CuteFnSpec { - canonical: "cute_rs::epilogue::__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::tile::load_tile", - func: CuteFn::LoadTile, - schema: &[SubstKind::Type, SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::tile::store_tile", - func: CuteFn::StoreTile, - schema: &[SubstKind::Type, SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::tile::store_global_no_allocate", - func: CuteFn::StoreGlobalNoAllocate, - schema: &[SubstKind::Type, SubstKind::Const], - }, - CuteFnSpec { - canonical: "cute_rs::tile::assume_div", - 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::tma::copy_tma_2d", - func: CuteFn::CopyTma2d, - schema: &[SubstKind::Type, SubstKind::Type], - }, - CuteFnSpec { - canonical: "cute_rs::tma::copy_tma_s2g_2d", - func: CuteFn::CopyTmaS2g2d, - schema: &[SubstKind::Type, SubstKind::Type], - }, - CuteFnSpec { - canonical: "cute_rs::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, - ], - }, -]; - -enum CuteIdentity { - /// Callee is not from the cute_rs crate; normal translation proceeds. - NotCuteCrate, - /// A recognized entry point. - Known(&'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(String), -} - -/// Definition modules containing ordinary Rust source abstractions rather -/// than importer-owned recognition stubs. -/// -/// This is deliberately a closed list. Adding a module means every function -/// defined directly or in an impl below it must have a real, device-safe body. -/// Panic-only recognition stubs are forbidden in these modules: they would be -/// body-translated instead of dispatched through [`CUTE_FNS`]. New modules -/// therefore require an explicit audit and an update to the exact-list test. -const CUTE_SOURCE_ABSTRACTION_MODULES: &[&str] = &[ - "arch", - "block_scaled", - "block_scaled_mma", - "epilogue", - "numeric", - "pipeline", - "scheduler", - "sm100", - "sm100_attention", - "tensor", - "tiled_copy", -]; - -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; - }; - let first_module = rest.split("::").next().unwrap_or_default(); - CUTE_SOURCE_ABSTRACTION_MODULES.contains(&first_module) -} - -/// Rebuild the canonical definition path by walking DefId parents, -/// mirroring `intrinsics::generated::classify_raw_intrinsic`. -fn canonical_path(fn_def: &FnDef, crate_name: &str) -> String { - let mut segments = Vec::new(); - let mut current = Some(fn_def.def_id()); - while let Some(def_id) = current { - let printed = def_id.name(); - let segment = printed.as_str().rsplit("::").next().unwrap_or_default(); - segments.push(segment.to_owned()); - current = def_id.parent(); - } - canonical_path_from_leaf_segments(crate_name, segments) -} - -/// Join a DefId parent walk without erasing a same-named nested module. -pub(super) fn canonical_path_from_leaf_segments( - crate_name: &str, - mut segments: Vec, -) -> String { - segments.reverse(); - // Only the first segment is the crate root. A nested module may legally - // have the same spelling and must remain part of the identity: - // - // cute_rs::cute_rs::tile::load_tile != cute_rs::tile::load_tile - if segments - .first() - .is_some_and(|segment| segment == crate_name) - { - segments.remove(0); - } - format!("{crate_name}::{}", segments.join("::")) -} - -#[cfg(test)] -mod canonical_path_tests { - use super::{ - block_scaled_method_name_from_path, canonical_path_from_leaf_segments, - is_source_abstraction_path, pipeline_method_name_from_path, - scheduler_method_name_from_path, smem_mma_method_name_from_path, - tensor_method_name_from_path, - }; - - #[test] - fn removes_only_the_crate_root_segment() { - let path = canonical_path_from_leaf_segments( - "cute_rs", - ["load_tile", "tile", "cute_rs", "cute_rs"] - .into_iter() - .map(str::to_string) - .collect(), - ); - assert_eq!(path, "cute_rs::cute_rs::tile::load_tile"); - } - - #[test] - fn cache_aware_global_store_has_an_exact_compiler_boundary() { - let spec = super::CUTE_FNS - .iter() - .find(|spec| spec.canonical == "cute_rs::tile::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] - ); - } - - #[test] - fn source_abstraction_modules_are_exact_and_closed() { - assert!(is_source_abstraction_path( - "cute_rs::arch::assume_warp_uniform", - "cute_rs" - )); - assert_eq!( - super::CUTE_SOURCE_ABSTRACTION_MODULES, - &[ - "arch", - "block_scaled", - "block_scaled_mma", - "epilogue", - "numeric", - "pipeline", - "scheduler", - "sm100", - "sm100_attention", - "tensor", - "tiled_copy", - ] - ); - assert!(is_source_abstraction_path( - "cute_rs::block_scaled::BlockScaledTensor<...>::thread_row", - "cute_rs" - )); - assert!(is_source_abstraction_path( - "cute-rs::numeric::PackedE2M1x2::to_f32x2", - "cute-rs" - )); - assert!(is_source_abstraction_path( - "cute_rs::tensor::Tensor<...>::layout", - "cute_rs" - )); - assert!(is_source_abstraction_path( - "cute_rs::tiled_copy::TiledCopy<...>::copy", - "cute_rs" - )); - assert!(is_source_abstraction_path( - "cute_rs::block_scaled_mma::Mxfp4TiledMma<...>::gemm", - "cute_rs" - )); - assert!(is_source_abstraction_path( - "cute_rs::epilogue::Sm120EpilogueWarp128x128::store_atom", - "cute_rs" - )); - assert!(is_source_abstraction_path( - "cute_rs::pipeline::TmaLoadPipeline<...>::consumer_wait", - "cute_rs" - )); - assert!(is_source_abstraction_path( - "cute_rs::scheduler::StaticPersistentTileScheduler<...>::current", - "cute_rs" - )); - assert!(!is_source_abstraction_path( - "cute_rs::tile::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" - )); - } - - #[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_linear"), - Some("store_linear") - ); - 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 - ); - } - - #[test] - fn block_scaled_method_path_ignores_only_the_impl_owner_spelling() { - assert_eq!( - block_scaled_method_name_from_path("cute_rs::block_scaled::KMajor>::from_slices"), - Some("from_slices") - ); - assert_eq!( - block_scaled_method_name_from_path("cute_rs::block_scaled::{impl#21}::dot_accumulate"), - Some("dot_accumulate") - ); - assert_eq!( - block_scaled_method_name_from_path( - "cute_rs::block_scaled::__compiler::block_scaled_dot_k64" - ), - Some("block_scaled_dot_k64"), - "the exact helper table takes precedence over method classification" - ); - assert_eq!( - block_scaled_method_name_from_path("cute_rs::other::{impl#1}::load"), - None - ); - } - - #[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}::coordinates"), - Some("coordinates") - ); - 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" - ); - } - - #[test] - fn pipeline_method_path_ignores_only_the_impl_owner_spelling() { - assert_eq!( - pipeline_method_name_from_path( - "cute_rs::pipeline::TmaLoadPipeline<...>>::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" - ); - } - - #[test] - fn shared_mma_method_path_keeps_module_and_method_exact() { - assert_eq!( - smem_mma_method_name_from_path( - "cute_rs::block_scaled_mma::Mxfp4TiledMma<...>>::load_a_128", - "block_scaled_mma" - ), - Some("load_a_128") - ); - assert_eq!( - smem_mma_method_name_from_path( - "cute-rs::tiled_copy::{impl#7}::from_raw_parts", - "tiled_copy" - ), - Some("from_raw_parts") - ); - assert_eq!( - smem_mma_method_name_from_path( - "cute_rs::other::Mxfp4TiledMma<...>>::load_a_128", - "block_scaled_mma" - ), - None - ); - assert_eq!( - smem_mma_method_name_from_path( - "cute_rs::block_scaled_mma::load_a_128", - "block_scaled_mma" - ), - None, - "a free function is not an inherent method" - ); - } - - #[test] - fn epilogue_method_paths_keep_each_owner_module_exact() { - assert_eq!( - smem_mma_method_name_from_path( - "cute_rs::epilogue::Sm120Epilogue128x128::sync_ready_for_tma", - "epilogue" - ), - Some("sync_ready_for_tma") - ); - assert_eq!( - smem_mma_method_name_from_path( - "cute_rs::block_scaled_mma::{impl#9}::store_tile", - "block_scaled_mma" - ), - Some("store_tile") - ); - assert_eq!( - pipeline_method_name_from_path( - "cute_rs::pipeline::TmaStorePipeline<1>::producer_commit" - ), - Some("producer_commit") - ); - assert_eq!( - smem_mma_method_name_from_path( - "cute_rs::epilogue::Sm120EpilogueWarp128x128::store_tile", - "block_scaled_mma" - ), - None, - "store_tile is defined in block_scaled_mma, not epilogue" - ); - } -} - -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), - } -} - -/// Recognize only the ergonomic inherent methods that expose the v0 tensor -/// flow. The DefId's direct parent is an implementation block, whose printed -/// number is intentionally ignored; the receiver/result ADT supplies the -/// stable identity instead. -fn tensor_method_name_from_path(path: &str) -> Option<&str> { - let mut segments = path.split("::"); - let crate_name = segments.next()?; - let module = segments.next()?; - // rustc_public currently prints an impl parent using the tail of its self - // type (for example `Tile1D>`), while older revisions exposed an - // `{impl#N}` segment. That spelling is not identity. The exact receiver - // and result ADTs checked by `classify_tensor_method_call` are. - let impl_owner = segments.next()?; - let method = segments.next()?; - if segments.next().is_some() - || !matches!(crate_name, "cute_rs" | "cute-rs") - || module != "tensor" - || impl_owner.is_empty() - { - return None; - } - Some(method) -} - -/// Return an inherent method name from the exact block-scaled module. -/// -/// rustc's spelling for the implementation owner is not stable. Receiver and -/// result ADTs below provide the identity instead. -fn block_scaled_method_name_from_path(path: &str) -> Option<&str> { - let mut segments = path.split("::"); - let crate_name = segments.next()?; - let module = segments.next()?; - let impl_owner = segments.next()?; - let method = segments.next()?; - if segments.next().is_some() - || !matches!(crate_name, "cute_rs" | "cute-rs") - || module != "block_scaled" - || impl_owner.is_empty() - { - return None; - } - Some(method) -} - -/// Return an inherent method name from the exact shared-copy or MMA module. -fn smem_mma_method_name_from_path<'a>(path: &'a str, expected_module: &str) -> Option<&'a str> { - let mut segments = path.split("::"); - let crate_name = segments.next()?; - let module = segments.next()?; - let impl_owner = segments.next()?; - let method = segments.next()?; - if segments.next().is_some() - || !matches!(crate_name, "cute_rs" | "cute-rs") - || module != expected_module - || impl_owner.is_empty() - { - return None; - } - Some(method) -} - -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, "tiled_copy") == 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, "block_scaled_mma")?; - 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) - } - "accumulate_k64" - if args.len() == 5 - && operand_ty(0) - .is_some_and(|ty| smem_mma::is_accumulator(&ty).unwrap_or(false)) => - { - Some(CuteFn::TiledGemm) - } - _ => None, - } -} - -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") { - 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, "block_scaled_mma") == 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. -fn scheduler_method_name_from_path(path: &str) -> Option<&str> { - let mut segments = path.split("::"); - let crate_name = segments.next()?; - let module = segments.next()?; - let impl_owner = segments.next()?; - let method = segments.next()?; - if segments.next().is_some() - || !matches!(crate_name, "cute_rs" | "cute-rs") - || module != "scheduler" - || impl_owner.is_empty() - { - return None; - } - Some(method) -} - -/// 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. -fn pipeline_method_name_from_path(path: &str) -> Option<&str> { - let mut segments = path.split("::"); - let crate_name = segments.next()?; - let module = segments.next()?; - let impl_owner = segments.next()?; - let method = segments.next()?; - if segments.next().is_some() - || !matches!(crate_name, "cute_rs" | "cute-rs") - || module != "pipeline" - || impl_owner.is_empty() - { - return None; - } - Some(method) -} - -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), - "slot" 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, - } -} - -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) - } - "current_tile" - if args.len() == 1 - && matches!((receiver_scheduler, destination_tile), (Some(source), Some(result)) if source == result) => - { - Some(CuteFn::SchedulerCurrentTile) - } - "coordinates" if args.len() == 1 && receiver_tile.is_some() => { - Some(CuteFn::WorkTileCoordinates) - } - "advance" if args.len() == 1 && receiver_scheduler.is_some() => { - Some(CuteFn::SchedulerAdvance) - } - _ => None, - } -} - -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, - } -} - -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, - }, - "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) - } - "store_linear" - if receiver_view.is_some_and(|view| { - view.access == CuteTensorAccessAttr::ReadWrite - && matches!(view.layout, CuteTensorLayoutAttr::Tile1D(_)) - }) => - { - Some(CuteFn::TensorStoreElementAbs) - } - _ => None, - } -} - -/// 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` when the callee is not a cute-rs function or is an explicit -/// source abstraction (the caller falls through to normal translation). -/// Returns `Some(result)` when this module owns the call, in which case the -/// shared path must not touch it. -#[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 be cute calls (indirect dispatch of cute_rs items is rejected by - // `requires_direct_dispatch` in the function-item path). - 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::emit(ctx, body, func, args, kind, target, block_ptr, prev_op, value_map, block_map, loc.clone()), - CuteFn::Sm100(kind) => sm100_emit::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::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( - "TmaLoadPipeline::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(), - ), - }) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[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 = canonical_path_from_leaf_segments( - "cute_rs", - [name, "mma", "cute_rs"] - .into_iter() - .map(str::to_string) - .collect(), - ); - 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}"); - } - } - #[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"); - } - } - } -} +pub(crate) use dispatch::try_translate_cute_call; +pub(super) use identity::canonical_path_from_leaf_segments; 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 similarity index 91% rename from cuda-oxide/crates/mir-importer/src/translator/cute/pipeline_emit.rs rename to cuda-oxide/crates/mir-importer/src/translator/cute/pipeline/emit.rs index 3d227828ec..4e9cbe84be 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/pipeline_emit.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/pipeline/emit.rs @@ -44,13 +44,15 @@ use pliron::r#type::{TypeHandle, Typed}; use pliron::value::Value; use rustc_public::mir; -use super::pipeline::{ +use super::super::pipeline::{ TmaLoadPipelineRust, 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, @@ -64,6 +66,7 @@ fn insert( } } +/// Branch from the last emitted effect to the call continuation; reject diverging calls. fn finish( ctx: &mut Context, effect: Ptr, @@ -81,6 +84,7 @@ fn finish( 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, @@ -108,6 +112,7 @@ fn finish_value( 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, @@ -127,6 +132,7 @@ fn destination_type( }) } +/// Build the concrete cursor carrier with its slot, phase and producer/consumer role. fn build_state( ctx: &mut Context, ty: TypeHandle, @@ -159,6 +165,7 @@ fn build_state( (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, @@ -182,6 +189,7 @@ fn pipeline_facts_from_type( }) } +/// Recover a cursor operand's stage count and producer/consumer role. fn state_facts_from_operand( operand: &mir::Operand, body: &mir::Body, @@ -205,6 +213,7 @@ fn state_facts_from_operand( }) } +/// Require the cursor to match the selected pipeline and expected participant role. fn require_state( state: CutePipelineStateAttr, stages: u64, @@ -229,6 +238,7 @@ fn require_state( Ok(()) } +/// Recover a shared barrier pointer at the validated unsafe pipeline boundary. fn recover_shared_barrier_base( ctx: &mut Context, base: Value, @@ -289,6 +299,7 @@ fn recover_shared_barrier_base( 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, @@ -328,6 +339,7 @@ fn field_addr( 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, @@ -350,14 +362,21 @@ fn load( (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, @@ -422,6 +441,7 @@ fn state_values( }) } +/// 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, @@ -478,6 +498,7 @@ pub(crate) fn emit_make( 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, @@ -517,6 +538,7 @@ pub(crate) fn emit_init( 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, @@ -582,6 +604,7 @@ pub(crate) fn emit_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, @@ -630,6 +653,7 @@ pub(crate) fn emit_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, @@ -709,14 +733,21 @@ pub(crate) fn emit_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, @@ -840,8 +871,10 @@ fn emit_lifecycle( } } +/// 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, @@ -866,6 +899,7 @@ 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, diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/pipeline.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/pipeline/mod.rs similarity index 85% rename from cuda-oxide/crates/mir-importer/src/translator/cute/pipeline.rs rename to cuda-oxide/crates/mir-importer/src/translator/cute/pipeline/mod.rs index 9bfac9ada6..d31b7c013e 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/pipeline.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/pipeline/mod.rs @@ -13,6 +13,8 @@ //! PipelineState -> { slot: u32, phase: u32 } //! ``` +pub(crate) mod emit; + use crate::error::{TranslationErr, TranslationResult}; use dialect_cute::attributes::CutePipelineStateAttr; use dialect_cute::types::CuteTmaLoadPipelineType; @@ -23,19 +25,27 @@ 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::TmaLoadPipeline`; public reexports do not change it. const TMA_LOAD_PIPELINE: &str = "cute_rs::pipeline::TmaLoadPipeline"; +/// 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 `TmaLoadPipeline`. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct TmaLoadPipelineRust { + /// 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); @@ -76,6 +86,7 @@ pub(crate) fn decode_tma_load_pipeline(ty: &Ty) -> Result Result, String> { let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { return Ok(None); @@ -151,6 +163,7 @@ pub(crate) fn decode_pipeline_state_receiver( } } +/// Construct the IR pipeline type with every validated static barrier parameter. fn pipeline_type(ctx: &mut Context, facts: TmaLoadPipelineRust) -> TypeHandle { CuteTmaLoadPipelineType::get( ctx, @@ -175,10 +188,12 @@ pub(crate) fn try_translate_tma_load_pipeline_type( } } +/// 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(); @@ -197,6 +212,7 @@ mod tests { 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()); 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..a5f3b0be60 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/registry.rs @@ -0,0 +1,662 @@ +/* + * 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, + /// 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, +} + +/// 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::sm100_attention::__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::sm100_attention::__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::sm100_attention::__compiler::tma_load", + func: CuteFn::Attention(sm100_attention::AttentionFn::TmaLoad), + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100_attention::__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::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::tiled_copy::__compiler::shared_tensor_overlay", + func: CuteFn::SharedTensorOverlay, + schema: &[ + SubstKind::Type, + SubstKind::Type, + SubstKind::Type, + SubstKind::Type, + ], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::tiled_mma_slice", + func: CuteFn::TiledMmaSlice, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::fragment_fill", + func: CuteFn::MmaFragmentFill, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::mma_load_scales", + func: CuteFn::MmaLoadScales, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::fragment_slice_k", + func: CuteFn::MmaFragmentSliceK, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::mma_load_a", + func: CuteFn::MmaLoadA, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::mma_partition_b", + func: CuteFn::MmaPartitionB, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::tiled_gemm", + func: CuteFn::TiledGemm, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::__compiler::epilogue_smem_overlay", + func: CuteFn::EpilogueOverlay, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::__compiler::epilogue_warp_slice", + func: CuteFn::EpilogueWarpSlice, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled_mma::__compiler::epilogue_store_fragment", + func: CuteFn::EpilogueStoreFragment, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::__compiler::epilogue_sync_reusable", + func: CuteFn::EpilogueSyncReusable, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::__compiler::epilogue_sync_ready_for_tma", + func: CuteFn::EpilogueSyncReadyForTma, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::__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::tile::load_tile", + func: CuteFn::LoadTile, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tile::store_tile", + func: CuteFn::StoreTile, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tile::store_global_no_allocate", + func: CuteFn::StoreGlobalNoAllocate, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tile::assume_div", + 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::tma::copy_tma_2d", + func: CuteFn::CopyTma2d, + schema: &[SubstKind::Type, SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::tma::copy_tma_s2g_2d", + func: CuteFn::CopyTmaS2g2d, + schema: &[SubstKind::Type, SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::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::super::identity::canonical_path_from_leaf_segments; + 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 = canonical_path_from_leaf_segments( + "cute_rs", + [name, "mma", "cute_rs"] + .into_iter() + .map(str::to_string) + .collect(), + ); + 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 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 similarity index 94% rename from cuda-oxide/crates/mir-importer/src/translator/cute/scheduler_emit.rs rename to cuda-oxide/crates/mir-importer/src/translator/cute/scheduler/emit.rs index e1cd3d1e44..621192aa9b 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/scheduler_emit.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/scheduler/emit.rs @@ -36,10 +36,11 @@ use pliron::r#type::{TypeHandle, Typed}; use pliron::value::Value; use rustc_public::mir; -use super::scheduler::{ +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, @@ -53,6 +54,7 @@ fn insert( } } +/// Branch from the last emitted effect to the call continuation; reject diverging calls. fn finish( ctx: &mut Context, effect: Ptr, @@ -70,6 +72,7 @@ fn finish( 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, @@ -97,6 +100,7 @@ fn finish_value( 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, @@ -117,6 +121,7 @@ fn destination_type( }) } +/// Construct a translated scheduler/work-tile carrier from its scalar fields. fn build_aggregate( ctx: &mut Context, aggregate_type: TypeHandle, @@ -149,6 +154,7 @@ fn build_aggregate( (aggregate, anchor) } +/// Decode the scheduler or work-tile receiver's compile-time grid. fn receiver_grid( body: &mir::Body, operand: &mir::Operand, @@ -179,6 +185,7 @@ fn receiver_grid( }) } +/// Find the mutable persistent-position field of a referenced scheduler. fn scheduler_field_addr( ctx: &mut Context, scheduler: Value, @@ -221,6 +228,7 @@ fn scheduler_field_addr( 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, @@ -243,6 +251,7 @@ fn load_field( (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, @@ -308,6 +317,7 @@ pub(crate) fn emit_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, @@ -358,6 +368,7 @@ pub(crate) fn emit_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, @@ -426,6 +437,7 @@ pub(crate) fn emit_current_tile( ) } +/// Project a work tile into its M, N and batch coordinates. #[allow(clippy::too_many_arguments)] pub(crate) fn emit_coordinates( ctx: &mut Context, @@ -480,6 +492,7 @@ pub(crate) fn emit_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, diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/scheduler.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/scheduler/mod.rs similarity index 83% rename from cuda-oxide/crates/mir-importer/src/translator/cute/scheduler.rs rename to cuda-oxide/crates/mir-importer/src/translator/cute/scheduler/mod.rs index 84017bde1c..ad37a7984c 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/scheduler.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/scheduler/mod.rs @@ -15,6 +15,8 @@ //! WorkTile !cute.work_tile> //! ``` +pub(crate) mod emit; + use crate::error::{TranslationErr, TranslationResult}; use dialect_cute::attributes::CuteTileGridAttr; use dialect_cute::types::CuteWorkTileType; @@ -25,9 +27,12 @@ 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::WorkTile`; public reexports do not change it. const WORK_TILE: &str = "cute_rs::scheduler::WorkTile"; +/// 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); @@ -39,6 +44,7 @@ fn exact_adt(ty: &Ty, expected: &str) -> Result, String> { 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")?; @@ -67,6 +73,7 @@ pub(crate) fn decode_work_tile(ty: &Ty) -> Result, Stri .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>, @@ -79,14 +86,17 @@ fn decode_behind_one_pointer( } } +/// 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() } @@ -105,10 +115,12 @@ pub(crate) fn try_translate_work_tile_type( } } +/// 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(); diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/sm100_emit.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/sm100.rs similarity index 85% rename from cuda-oxide/crates/mir-importer/src/translator/cute/sm100_emit.rs rename to cuda-oxide/crates/mir-importer/src/translator/cute/sm100.rs index 1045902438..8efdfc5e3a 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/sm100_emit.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/sm100.rs @@ -27,37 +27,59 @@ use pliron::value::Value; use rustc_public::mir; use rustc_public::ty::{GenericArgKind, GenericArgs, RigidTy, TyKind}; -use super::smem_mma_emit::{finish, insert}; +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")); @@ -68,6 +90,7 @@ fn generics(func: &mir::Operand) -> TranslationResult { 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!( @@ -78,6 +101,7 @@ fn constant(args: &GenericArgs, index: usize, loc: &Location) -> TranslationResu 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!( @@ -89,6 +113,7 @@ fn layout(args: &GenericArgs, index: usize) -> TranslationResult &'static [usize] { use Sm100Fn::*; match kind { @@ -135,6 +162,7 @@ fn shared_operands(kind: Sm100Fn) -> &'static [usize] { } } +/// Return the exact scalar-carrier count required by this compiler hook. fn operand_count(kind: Sm100Fn) -> usize { use Sm100Fn::*; match kind { @@ -149,6 +177,7 @@ fn operand_count(kind: Sm100Fn) -> usize { } } +/// 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, diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/sm100_attention_emit.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/sm100_attention.rs similarity index 89% rename from cuda-oxide/crates/mir-importer/src/translator/cute/sm100_attention_emit.rs rename to cuda-oxide/crates/mir-importer/src/translator/cute/sm100_attention.rs index 7811106c57..e31f740b71 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/sm100_attention_emit.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/sm100_attention.rs @@ -27,21 +27,28 @@ use pliron::value::Value; use rustc_public::mir; use rustc_public::ty::{GenericArgKind, GenericArgs, RigidTy, TyKind}; -use super::smem_mma_emit::{finish, insert}; +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")); @@ -52,6 +59,7 @@ fn generics(func: &mir::Operand) -> TranslationResult { 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!( @@ -62,6 +70,7 @@ fn constant(args: &GenericArgs, index: usize, loc: &Location) -> TranslationResu 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")); @@ -70,6 +79,7 @@ fn layout(args: &GenericArgs, index: usize) -> TranslationResult<(CuteComposedLa .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], @@ -78,6 +88,7 @@ fn shared_operands(kind: AttentionFn) -> &'static [usize] { 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, @@ -86,6 +97,7 @@ fn operand_count(kind: AttentionFn) -> usize { 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, 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 similarity index 93% rename from cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma_emit.rs rename to cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma/emit.rs index 04070c5806..754bed394b 100644 --- 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 @@ -43,14 +43,15 @@ use pliron::r#type::{TypeHandle, Typed}; use pliron::value::Value; use rustc_public::mir; -use super::emit::{load_through, recover_cta_shared_pointer, struct_field_addr}; -use super::smem_mma::{ +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, }; -pub(super) fn insert( +/// 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, @@ -63,7 +64,8 @@ pub(super) fn insert( } } -pub(super) fn finish( +/// 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, @@ -80,8 +82,9 @@ pub(super) fn finish( 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(super) fn finish_value( +pub(in crate::translator::cute) fn finish_value( ctx: &mut Context, producer: Ptr, value: Value, @@ -107,7 +110,8 @@ pub(super) fn finish_value( finish(ctx, after, target, block_map, loc, what) } -pub(super) fn destination_rust_type( +/// 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, @@ -119,7 +123,8 @@ pub(super) fn destination_rust_type( }) } -pub(super) fn destination_type( +/// 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, @@ -134,7 +139,8 @@ pub(super) fn destination_type( }) } -pub(super) fn operand_rust_type( +/// 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, @@ -146,7 +152,8 @@ pub(super) fn operand_rust_type( }) } -pub(super) fn aggregate_field( +/// 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, @@ -195,7 +202,8 @@ pub(super) fn aggregate_field( Ok((op.deref(ctx).get_result(0), op)) } -pub(super) fn build_struct_prefix( +/// 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], @@ -250,7 +258,8 @@ pub(super) fn build_struct_prefix( Ok((value, anchor)) } -pub(super) fn pointer_cast( +/// 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, @@ -299,14 +308,17 @@ pub(super) fn pointer_cast( 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 { @@ -341,6 +353,7 @@ fn smem_view_type(ctx: &mut Context, facts: &SharedTensorRust) -> (TypeHandle, T ) } +/// Require a recognized shared tensor operand and recover its layout and role. fn decode_shared_operand( body: &mir::Body, operand: &mir::Operand, @@ -360,6 +373,7 @@ fn decode_shared_operand( }) } +/// Require a recognized tiled-MMA operand and recover its shared placement. fn decode_mma_operand( body: &mir::Body, operand: &mir::Operand, @@ -379,6 +393,7 @@ fn decode_mma_operand( }) } +/// Extract a shared tensor's storage pointer and capacity from its Rust carrier. #[allow(clippy::too_many_arguments)] fn shared_carriers( ctx: &mut Context, @@ -401,6 +416,7 @@ fn shared_carriers( 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, @@ -415,6 +431,7 @@ fn mma_lane( 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, @@ -436,7 +453,8 @@ fn emit_zero_f32( (op.deref(ctx).get_result(0), op) } -pub(super) fn struct_field_type( +/// 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, @@ -452,6 +470,7 @@ pub(super) fn struct_field_type( }) } +/// 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, @@ -543,6 +562,7 @@ pub(crate) fn emit_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, @@ -603,6 +623,7 @@ pub(crate) fn emit_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, @@ -671,6 +692,7 @@ pub(crate) fn emit_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, @@ -786,6 +808,7 @@ pub(crate) fn emit_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, @@ -850,6 +873,7 @@ pub(crate) fn emit_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, @@ -935,6 +959,7 @@ pub(crate) fn emit_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, @@ -1044,6 +1069,7 @@ pub(crate) fn emit_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, diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma/mod.rs similarity index 62% rename from cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma.rs rename to cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma/mod.rs index 942c167913..97d44092e2 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma/mod.rs @@ -8,6 +8,8 @@ //! 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}; @@ -15,36 +17,57 @@ 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::tiled_copy::SharedTensor`; public reexports do not change it. const SHARED_TENSOR: &str = "cute_rs::tiled_copy::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"; -const SM120_SCALE_ATOM: &str = "cute_rs::block_scaled::Sm120ScaleAtom"; +/// 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"; -const TILED_MMA: &str = "cute_rs::block_scaled_mma::Mxfp4TiledMma"; -const ACCUMULATOR: &str = "cute_rs::block_scaled_mma::Mxf4AccumulatorTile2x8"; -const SCALE_STAGE: &str = "cute_rs::block_scaled_mma::Mxf4ScaleTile128"; -const SCALE_K64: &str = "cute_rs::block_scaled_mma::Mxf4ScalePairs128"; -const B_TILE_K64: &str = "cute_rs::block_scaled_mma::Mxf4BTileK64"; - +/// Exact defining path of `cute_rs::block_scaled_mma::tiled::Mxfp4TiledMma`; public reexports do not change it. +const TILED_MMA: &str = "cute_rs::block_scaled_mma::tiled::Mxfp4TiledMma"; +/// Exact defining path of `cute_rs::block_scaled_mma::fragments::Mxf4AccumulatorTile2x8`; public reexports do not change it. +const ACCUMULATOR: &str = "cute_rs::block_scaled_mma::fragments::Mxf4AccumulatorTile2x8"; +/// Exact defining path of `cute_rs::block_scaled_mma::fragments::Mxf4ScaleTile128`; public reexports do not change it. +const SCALE_STAGE: &str = "cute_rs::block_scaled_mma::fragments::Mxf4ScaleTile128"; +/// Exact defining path of `cute_rs::block_scaled_mma::fragments::Mxf4ScalePairs128`; public reexports do not change it. +const SCALE_K64: &str = "cute_rs::block_scaled_mma::fragments::Mxf4ScalePairs128"; +/// Exact defining path of `cute_rs::block_scaled_mma::fragments::Mxf4BTileK64`; public reexports do not change it. +const B_TILE_K64: &str = "cute_rs::block_scaled_mma::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; @@ -52,6 +75,7 @@ fn adt(ty: &Ty) -> Option<(String, GenericArgs)> { 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); @@ -63,6 +87,7 @@ fn exact_marker(ty: &Ty, expected: &str) -> Result { 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) @@ -73,14 +98,17 @@ fn role(ty: &Ty) -> Result { } } +/// 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 inner = "(32,4):(4,1)" .parse() @@ -88,6 +116,7 @@ fn scale_placement() -> ComposedLayout { ComposedLayout::from_layout(inner, OffsetUnit::Elements) } +/// Return the fixed swizzled 128-by-32 packed-data placement. pub(crate) fn canonical_data_placement() -> ComposedLayout { let inner: Layout = "(128,32):(32,1)" .parse() @@ -96,6 +125,7 @@ pub(crate) fn canonical_data_placement() -> ComposedLayout { .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); @@ -129,6 +159,7 @@ pub(crate) fn decode_shared_tensor(ty: &Ty) -> Result, 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>, @@ -141,10 +172,12 @@ fn behind_one_pointer( } } +/// 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); @@ -158,27 +191,33 @@ pub(crate) fn decode_tiled_mma(ty: &Ty) -> Result, String> })) } +/// 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); @@ -192,15 +231,18 @@ pub(crate) fn decode_b_tile(ty: &Ty) -> Result, String> { })) } +/// 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().inner().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(); 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 index cbd1d0f437..fad8ae1ad7 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/static_config.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/static_config.rs @@ -28,27 +28,44 @@ use rustc_public::ty::{ 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::T2`; public reexports do not change it. const T2: &str = "cute_rs::markers::T2"; +/// Exact defining path of `cute_rs::markers::L`; public reexports do not change it. const LAYOUT: &str = "cute_rs::markers::L"; +/// 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::CpAsync`; public reexports do not change it. const CP_ASYNC: &str = "cute_rs::markers::CpAsync"; +/// 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::Composed`; public reexports do not change it. const COMPOSED: &str = "cute_rs::markers::Composed"; +/// 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, @@ -56,13 +73,16 @@ pub(crate) struct CooperativeConfig { #[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 { @@ -100,6 +120,7 @@ fn cooperative_element(ty: &Ty) -> Result<(i64, &'static str), String> { Ok(info) } +/// Recover an ADT's defining path by walking its compiler identity, not its display name. fn canonical_adt_path(def: &AdtDef) -> String { let crate_name = def.krate().name.to_string(); let mut segments = Vec::new(); @@ -113,6 +134,7 @@ fn canonical_adt_path(def: &AdtDef) -> String { super::canonical_path_from_leaf_segments(&crate_name, segments) } +/// 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:?}")); @@ -124,6 +146,7 @@ fn adt_owned(ty: &Ty, expected: &str) -> Result { 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 @@ -145,6 +168,7 @@ fn expect_schema(args: &GenericArgs, expected: &[&str], owner: &str) -> Result<( } } +/// 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 @@ -159,6 +183,7 @@ fn unsigned_const(value: &TyConst, what: &str) -> Result { 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 @@ -173,6 +198,7 @@ fn signed_const(value: &TyConst, what: &str) -> Result { 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), @@ -180,6 +206,7 @@ fn type_arg(args: &GenericArgs, index: usize, owner: &str) -> Result } } +/// 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()), @@ -187,6 +214,7 @@ fn const_arg(args: &GenericArgs, index: usize, owner: &str) -> Result Result { let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { return Err(format!("layout tuple marker must be an ADT, found {ty:?}")); @@ -210,6 +238,7 @@ fn decode_tuple(ty: &Ty) -> Result { } } +/// 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), @@ -221,6 +250,7 @@ fn checked_shape_size(shape: &IntTuple) -> Result { } } +/// 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:?}")); @@ -257,6 +287,7 @@ pub(crate) fn decode_layout(ty: &Ty) -> Result { Ok(layout) } +/// Decode and validate a `CpAsync` copy atom's byte width. fn decode_atom(ty: &Ty) -> Result { let args = adt_owned(ty, CP_ASYNC)?; expect_schema(&args, &["Const"], CP_ASYNC)?; @@ -273,6 +304,7 @@ fn decode_atom(ty: &Ty) -> Result { 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)?; @@ -286,6 +318,7 @@ fn decode_leading_dim_divisor(ty: &Ty) -> Result { 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)?; @@ -302,7 +335,8 @@ fn decode_swizzle(ty: &Ty) -> Result { /// Decode the sealed attention layout identity and its major mode. pub(crate) fn decode_attention_layout(ty: &Ty) -> Result<(ComposedLayout, bool), String> { - const PATH: &str = "cute_rs::sm100_attention::SharedLayout"; + /// Exact defining path of `cute_rs::sm100_attention::shared::SharedLayout`; public reexports do not change it. + const PATH: &str = "cute_rs::sm100_attention::shared::SharedLayout"; let args = adt_owned(ty, PATH)?; expect_schema(&args, &["Const", "Const", "Const"], PATH)?; let rows = u32::try_from(unsigned_const( @@ -326,6 +360,7 @@ pub(crate) fn decode_attention_layout(ty: &Ty) -> Result<(ComposedLayout, bool), )) } +/// 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 @@ -446,7 +481,9 @@ pub(crate) fn decode_load_matrix_callee(func: &mir::Operand) -> Result Result, @@ -48,6 +49,7 @@ fn insert( } } +/// Branch from the last emitted effect to the call continuation; reject diverging calls. fn finish( ctx: &mut Context, effect: Ptr, @@ -65,6 +67,7 @@ fn finish( 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, @@ -92,6 +95,7 @@ fn finish_value( 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, @@ -115,6 +119,7 @@ fn extract_field( (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, @@ -155,6 +160,7 @@ fn load_pointer_pointee( 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, @@ -182,6 +188,7 @@ fn load_tensor_if_referenced( 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, @@ -223,6 +230,7 @@ fn cast_carrier_to_element( 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, @@ -256,6 +264,7 @@ fn view_transfer_facts( 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, @@ -397,6 +406,7 @@ pub(crate) fn emit_make( ) } +/// 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, @@ -458,6 +468,7 @@ pub(crate) fn emit_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, @@ -500,6 +511,7 @@ pub(crate) fn emit_slice( ) } +/// 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, @@ -548,6 +560,7 @@ pub(crate) fn emit_query( ) } +/// Load a complete per-thread tile into a register carrier. #[allow(clippy::too_many_arguments)] pub(crate) fn emit_load( ctx: &mut Context, @@ -611,6 +624,7 @@ pub(crate) fn emit_load( 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, @@ -675,6 +689,7 @@ pub(crate) fn emit_store( 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, diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/tensor.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/tensor/mod.rs similarity index 84% rename from cuda-oxide/crates/mir-importer/src/translator/cute/tensor.rs rename to cuda-oxide/crates/mir-importer/src/translator/cute/tensor/mod.rs index eb4f884eba..ab31fe4da6 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/cute/tensor.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/tensor/mod.rs @@ -17,6 +17,8 @@ //! `!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::{ @@ -31,29 +33,42 @@ 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 defining path by walking its compiler identity, not its display name. pub(super) fn canonical_adt_path(def: &AdtDef) -> String { let crate_name = def.krate().name.to_string(); let mut segments = Vec::new(); @@ -67,6 +82,7 @@ pub(super) fn canonical_adt_path(def: &AdtDef) -> String { super::canonical_path_from_leaf_segments(&crate_name, segments) } +/// Require the expected ordered generic-argument kinds for a source tensor type. pub(super) fn exact_schema( args: &GenericArgs, expected: &[&str], @@ -92,6 +108,7 @@ pub(super) fn exact_schema( } } +/// 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), @@ -99,6 +116,7 @@ pub(super) fn type_arg(args: &GenericArgs, index: usize, owner: &str) -> Result< } } +/// 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()), @@ -106,6 +124,7 @@ pub(super) fn const_arg(args: &GenericArgs, index: usize, owner: &str) -> Result } } +/// 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 @@ -120,6 +139,7 @@ pub(super) fn unsigned_const(value: &TyConst, what: &str) -> Result 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); 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..2b0f8097fb 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_alloc.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_alloc.rs @@ -506,10 +506,10 @@ fn decode_relocation_addend( /// Materialize a thin pointer field from an aggregate constant's allocation. /// /// Aggregate **const** values with thin pointers to device statics are -/// materialized via [`MirGlobalAllocOp`] per field (addend taken from the +/// materialized via [`MirGlobalAllocOp`](dialect_mir::ops::MirGlobalAllocOp) per field (addend taken from the /// relocation's stored bytes). This does **not** mean device-static /// *initializers* that themselves contain pointer relocations are supported — -/// those remain rejected by [`allocation_initializer_data`]. +/// those remain rejected by [`allocation_initializer_data`](super::statics::allocation_initializer_data). /// /// When the field has no provenance entry, falls back to the existing /// inttoptr-of-bytes path (null / exposed provenance). diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_bytes.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_bytes.rs index 8bc2ade2f8..1ee35ffa2f 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_bytes.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_bytes.rs @@ -37,7 +37,7 @@ use std::num::NonZeroUsize; /// [..]` indexed by runtime value) to a `MirConstructArrayOp`. Element stride /// and aggregate field offsets come from rustc layout. Thin pointer fields /// inside elements that relocate to device statics are materialized via -/// [`MirGlobalAllocOp`] per field. +/// [`MirGlobalAllocOp`](dialect_mir::ops::MirGlobalAllocOp) per field. pub(super) fn translate_array_value_constant( ctx: &mut Context, constant: &mir::ConstOperand, @@ -466,7 +466,7 @@ fn translate_array_value_constant_inner( /// constant is a promoted `&Struct`) /// 2. Decode each field at its rustc layout offset from that allocation /// 3. Thin pointer fields that relocate to device statics become -/// [`MirGlobalAllocOp`] results; other fields use the byte decoder +/// [`MirGlobalAllocOp`](dialect_mir::ops::MirGlobalAllocOp) results; other fields use the byte decoder /// 4. Create MirConstructStructOp with those operands /// /// Each field is read at the byte offset rustc's layout records for it, so diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_addr.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_addr.rs index 0b09daa48c..c7ef3e3158 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 @@ -570,11 +570,11 @@ pub(super) fn emit_array_subslice_value( /// Walks the projection chain and emits the correct pliron ops for each /// element: /// -/// - `Field(idx, _)` → [`MirFieldAddrOp`] +/// - `Field(idx, _)` → [`MirFieldAddrOp`](dialect_mir::ops::MirFieldAddrOp) /// - `ConstantIndex {offset, from_end, ..}` → an element address. Forward /// indexes materialize `offset` directly; from-end indexes on slices extract /// the fat-pointer length and materialize `len - offset` at runtime. -/// - `Index(local)` → `load_local(local)` + [`MirArrayElementAddrOp`] +/// - `Index(local)` → `load_local(local)` + [`MirArrayElementAddrOp`](dialect_mir::ops::MirArrayElementAddrOp) /// (array pointee) or `load_local(local)` + [`MirPtrOffsetOp`] (slice data pointer) /// - `Deref` → `MirLoadOp` of the pointer (the loaded pointer IS /// the pointee's address); subsequent projections apply to the pointee. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs b/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs index cd5fe4867a..44e1db553a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs @@ -337,20 +337,21 @@ fn is_cute_recognition_stub_identity( }; let module = match impl_owner { CuteImplOwner::Tensor | CuteImplOwner::TensorMut => "tensor", - CuteImplOwner::BlockScaledTensor - | CuteImplOwner::BlockScaledThreadRow + CuteImplOwner::BlockScaledTensor => "block_scaled::tensor", + CuteImplOwner::BlockScaledThreadRow | CuteImplOwner::BlockScaledTile64 - | CuteImplOwner::LoadedBlockScaledTile64 => "block_scaled", + | CuteImplOwner::LoadedBlockScaledTile64 => "block_scaled::gemv", CuteImplOwner::StaticPersistentTileScheduler | CuteImplOwner::WorkTile => "scheduler", CuteImplOwner::TmaLoadPipeline | CuteImplOwner::PipelineState | CuteImplOwner::TmaStorePipeline => "pipeline", CuteImplOwner::SharedTensor => "tiled_copy", - CuteImplOwner::Mxfp4TiledMma - | CuteImplOwner::Mxf4AccumulatorTile2x8 - | CuteImplOwner::Mxf4ScaleTile128 => "block_scaled_mma", + CuteImplOwner::Mxfp4TiledMma => "block_scaled_mma::tiled", + CuteImplOwner::Mxf4AccumulatorTile2x8 | CuteImplOwner::Mxf4ScaleTile128 => { + "block_scaled_mma::fragments" + } CuteImplOwner::Sm120Epilogue128x128 => "epilogue", - CuteImplOwner::Sm120EpilogueWarp128x128 => "block_scaled_mma", + CuteImplOwner::Sm120EpilogueWarp128x128 => "block_scaled_mma::epilogue", }; let Some(rest) = definition_path .strip_prefix("::") @@ -3146,17 +3147,17 @@ mod tests { )); assert!(is_cute_recognition_stub_identity( "cute_rs", - "::block_scaled::{impl#17}::from_slices", + "::block_scaled::tensor::{impl#17}::from_slices", Some(CuteImplOwner::BlockScaledTensor), )); assert!(is_cute_recognition_stub_identity( "cute_rs", - "::block_scaled::{impl#19}::k_tile", + "::block_scaled::gemv::{impl#19}::k_tile", Some(CuteImplOwner::BlockScaledThreadRow), )); assert!(is_cute_recognition_stub_identity( "cute_rs", - "::block_scaled::{impl#21}::dot_accumulate", + "::block_scaled::gemv::{impl#21}::dot_accumulate", Some(CuteImplOwner::LoadedBlockScaledTile64), )); assert!(is_cute_recognition_stub_identity( @@ -3214,17 +3215,17 @@ mod tests { )); assert!(is_cute_recognition_stub_identity( "cute_rs", - "::block_scaled_mma::{impl#7}::load_a_128", + "::block_scaled_mma::tiled::{impl#7}::load_a_128", Some(CuteImplOwner::Mxfp4TiledMma), )); assert!(is_cute_recognition_stub_identity( "cute_rs", - "::block_scaled_mma::{impl#4}::zero", + "::block_scaled_mma::fragments::{impl#4}::zero", Some(CuteImplOwner::Mxf4AccumulatorTile2x8), )); assert!(is_cute_recognition_stub_identity( "cute_rs", - "::block_scaled_mma::{impl#5}::pairs_at_unchecked", + "::block_scaled_mma::fragments::{impl#5}::pairs_at_unchecked", Some(CuteImplOwner::Mxf4ScaleTile128), )); assert!(is_cute_recognition_stub_identity( @@ -3239,7 +3240,7 @@ mod tests { )); assert!(is_cute_recognition_stub_identity( "cute_rs", - "::block_scaled_mma::{impl#9}::store_tile", + "::block_scaled_mma::epilogue::{impl#9}::store_tile", Some(CuteImplOwner::Sm120EpilogueWarp128x128), )); assert!(is_cute_recognition_stub_identity( @@ -3266,7 +3267,7 @@ mod tests { )); assert!(!is_cute_recognition_stub_identity( "cute_rs", - "::block_scaled::{impl#21}::value_pair", + "::block_scaled::gemv::{impl#21}::value_pair", Some(CuteImplOwner::LoadedBlockScaledTile64), )); assert!(!is_cute_recognition_stub_identity( @@ -3281,7 +3282,7 @@ mod tests { )); assert!(!is_cute_recognition_stub_identity( "cute_rs", - "::block_scaled_mma::{impl#7}::load_b_pair", + "::block_scaled_mma::tiled::{impl#7}::load_b_pair", Some(CuteImplOwner::Mxfp4TiledMma), )); assert!(!is_cute_recognition_stub_identity( @@ -3296,6 +3297,56 @@ mod tests { )); } + #[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", + ), + ( + "block_scaled_mma::tiled", + CuteImplOwner::Mxfp4TiledMma, + "load_a_128", + ), + ( + "block_scaled_mma::fragments", + CuteImplOwner::Mxf4AccumulatorTile2x8, + "zero", + ), + ( + "block_scaled_mma::epilogue", + CuteImplOwner::Sm120EpilogueWarp128x128, + "store_tile", + ), + ] { + 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/tests/fixtures/primitives_api_compile/src/kernel.rs b/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/src/kernel.rs index 52cc6cc42b..5c96c6e00e 100644 --- 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 @@ -1,6 +1,6 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use cuda_host::cuda_module; +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. @@ -12,10 +12,16 @@ fn generic_shared_identity(pointer: *mut u8) -> *mut u8 { #[cuda_module] mod probes { use ctm_rs::{ + async_copy, cluster::{self, CtaMask}, + convert, float, groups::CtaGroup2, + kernel, launch_bounds, launch_contract, mbarrier::{self, Barrier, BarrierAddressExt, ManagedBarrier, Ready, TmaBarrier, Uninit}, - memory::{F16Tile128, F16TileB80, F16TileB96, F16TileC, SharedRegion}, + memory::{ + DisjointSlice, DynamicSharedArray, F16Tile128, F16TileB80, F16TileB96, F16TileC, + SharedRegion, + }, tcgen05::{ self, descriptors::*, @@ -26,7 +32,7 @@ mod probes { tma::{self, TensorMap}, warp, }; - use cuda_device::{SharedArray, cluster_launch, kernel}; + use cuda_device::{SharedArray, cluster_launch}; static mut SLOT: SharedArray = SharedArray::UNINIT; static mut BAR: Barrier = Barrier::UNINIT; @@ -109,7 +115,7 @@ mod probes { #[cluster_launch(2, 1, 1)] pub fn primitives_sync(out: *mut u64) { unsafe { - let shared = cuda_device::DynamicSharedArray::::get_raw(); + let shared = DynamicSharedArray::::get_raw(); out.add(3) .write(super::generic_shared_identity(shared) as u64); mbarrier::fence_mbarrier_init_release_cluster(); @@ -138,6 +144,34 @@ mod probes { } } + // 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) { diff --git a/cuda-oxide/cuteir/README.md b/cuda-oxide/cuteir/README.md index 8575fc1373..1ff09d9492 100644 --- a/cuda-oxide/cuteir/README.md +++ b/cuda-oxide/cuteir/README.md @@ -77,6 +77,14 @@ with the same event-timing policy. Kernel names distinguish `_cute_rs` and commands. These use CuTe layouts plus NVVM operations; `cute_ext` GQA examples belong to the separate LIR programming model. +The RTX 5090 attention examples target **SM120a**. They use `cute_rs` for +layouts and copies, and `ctm_rs` for instruction-level operations: + +| Example | Purpose | +| :--- | :--- | +| [fmha_prefill_sm120a_cute](examples/fmha_prefill_sm120a_cute) | FP16 D64 prefill with causal or noncausal attention. | +| [fmha_decode_sm120a_cute](examples/fmha_decode_sm120a_cute) | Single-query FP16 D64 decode with partial-result reduction. | + The original examples are deliberately ordered from smallest to largest: ```text @@ -85,6 +93,37 @@ vector tiles → scaled row tiles → global/shared matrix tiles load + add load + dot TMA + pipeline + MMA + epilogue ``` +## Maintaining the CuTe implementation + +The Rust API, importer, dialect and exporter have different responsibilities: + +| Location | Responsibility | +| :--- | :--- | +| [cute-rs/src/](cute-rs/src) | Public kernel abstractions, with private implementation modules for layouts, fragments, MMA, copies and pipelines. Existing root and family imports remain the public API. | +| [cute-layout/src/validation/](cute-layout/src/validation) | Independent proofs for cooperative copies, thread/value coverage, matrix loads and TMA layouts. | +| [mir-importer/src/translator/cute/](../crates/mir-importer/src/translator/cute) | Exact compiler-entry registration, defining-path recognition and operation dispatch. Each family groups its type decoder and emitter. | +| [dialect-cute/src/](dialect-cute/src) | Backend-neutral operations, attributes and types. `verify/` separates shared control-flow and value-provenance analysis from operation-family protocol checks. | +| [cuda-oxide-mlir-export/src/cute/](../crates/cuda-oxide-mlir-export/src/cute) | Mapping recipes grouped by tensor, copy, MMA, pipeline, epilogue and scheduler, with architecture-specific implementations inside those families. | + +CuTe-RS's compiler boundary uses **defining paths**, not public reexport paths. +Moving a type or an inherent implementation into a child module changes its +compiler identity even when the old Rust import remains valid. Keep compiler +hooks in their existing `__compiler` modules. When moving a recognized type or +method, update the importer's exact type/layout decoders and method classifier, +and the backend collector's exact implementation-owner paths together. Preserve +the closed source-abstraction policy; accepting arbitrary suffix matches can +silently translate a compiler hook as ordinary Rust. + +Validate such changes with importer and collector identity tests, compile-fail +API tests, and real kernel compilation. Host tests alone do not exercise the +compiler's recognition of device calls. Compare emitted MLIR and instructions +for representative tensor, copy, GEMV, GEMM and attention kernels. + +Keep rustdoc short and close to the API: explain its purpose, supported shapes, +units and caller obligations. Add a small Rust example when it makes usage or +ordering clearer, and check it as a doctest where possible. Link shared details +instead of repeating them in each module. + ## Run the examples Run these commands from the `cuda-oxide/` directory after completing the normal diff --git a/cuda-oxide/cuteir/ctm-rs/src/async_copy.rs b/cuda-oxide/cuteir/ctm-rs/src/async_copy.rs new file mode 100644 index 0000000000..6fad77366d --- /dev/null +++ b/cuda-oxide/cuteir/ctm-rs/src/async_copy.rs @@ -0,0 +1,38 @@ +/* + * 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. + +/// 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) } +} diff --git a/cuda-oxide/cuteir/ctm-rs/src/cluster.rs b/cuda-oxide/cuteir/ctm-rs/src/cluster.rs index 9c992354d7..7f812f7f64 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/cluster.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/cluster.rs @@ -10,6 +10,7 @@ use super::{ 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), "`].")] @@ -100,7 +101,9 @@ pub unsafe fn 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 {} @@ -109,11 +112,14 @@ impl Clone for CtaMask { *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 { @@ -125,6 +131,7 @@ impl CtaMask { }) } } + /// Return the mask bits consumed by the multicast instruction. #[inline(always)] pub const fn bits(self) -> u16 { self.bits diff --git a/cuda-oxide/cuteir/ctm-rs/src/convert.rs b/cuda-oxide/cuteir/ctm-rs/src/convert.rs new file mode 100644 index 0000000000..2dcd42077e --- /dev/null +++ b/cuda-oxide/cuteir/ctm-rs/src/convert.rs @@ -0,0 +1,17 @@ +/* + * 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) +} diff --git a/cuda-oxide/cuteir/ctm-rs/src/float.rs b/cuda-oxide/cuteir/ctm-rs/src/float.rs new file mode 100644 index 0000000000..9f90c280b5 --- /dev/null +++ b/cuda-oxide/cuteir/ctm-rs/src/float.rs @@ -0,0 +1,29 @@ +/* + * 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) +} diff --git a/cuda-oxide/cuteir/ctm-rs/src/groups.rs b/cuda-oxide/cuteir/ctm-rs/src/groups.rs index 54380ea9e5..91ca3259ee 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/groups.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/groups.rs @@ -3,10 +3,14 @@ * 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 { @@ -19,6 +23,7 @@ pub trait CtaGroup: sealed::Group { /// 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 {} diff --git a/cuda-oxide/cuteir/ctm-rs/src/host/tma.rs b/cuda-oxide/cuteir/ctm-rs/src/host/tma.rs index c67fbbf73d..d5f4129d3b 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/host/tma.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/host/tma.rs @@ -10,11 +10,14 @@ 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 {} @@ -33,40 +36,60 @@ impl TmaElement for u16 { 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(driver::CUresult), + 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, @@ -117,6 +140,7 @@ fn validate_2d>( } // 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> { @@ -167,11 +191,13 @@ pub fn encode_2d>( }) } +/// 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}; @@ -224,6 +250,7 @@ mod tests { ); } + /// 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}; @@ -242,6 +269,7 @@ mod tests { } } + /// 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| { @@ -256,6 +284,7 @@ mod tests { )); } + /// Validated encoder parameters must preserve tile extents, swizzle spans and cache options. #[test] fn tma_arguments_preserve_physical_shapes_and_options() { let p = @@ -298,6 +327,7 @@ mod tests { ); } + /// 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 [ diff --git a/cuda-oxide/cuteir/ctm-rs/src/lib.rs b/cuda-oxide/cuteir/ctm-rs/src/lib.rs index 8f5aa26401..2f0d8c4e18 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/lib.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/lib.rs @@ -8,6 +8,20 @@ //! 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 and examples are in `cuda-oxide/cuteir/ctm-rs/README.md`; its +//! `docs/primitives-coverage.md` lists supported instruction contracts. +//! 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] @@ -18,7 +32,13 @@ 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::{cuda_module, kernel, launch_bounds, launch_contract}; + +pub mod async_copy; pub mod cluster; +pub mod convert; +pub mod float; pub mod groups; pub mod mbarrier; pub mod memory; @@ -33,6 +53,7 @@ pub mod warp; /// Hardware thread/block coordinates, without inferred thread mapping. pub mod thread { + /// 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), "`].")] diff --git a/cuda-oxide/cuteir/ctm-rs/src/memory.rs b/cuda-oxide/cuteir/ctm-rs/src/memory.rs index 02d7163dc4..12399cd24b 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/memory.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/memory.rs @@ -8,29 +8,58 @@ 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. +pub use cuda_device::{DisjointSlice, DynamicSharedArray}; + +/// 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 {} -pub struct Swizzled2D(PhantomData); +/// 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> {} @@ -77,23 +106,31 @@ impl PhysicalFormat for Mxf4ScaleTile { const SWIZZLE_BYTES: usize = 0; } +/// Byte displacement used when encoding a shared-memory instruction address. #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct ByteOffset(usize); +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> {} @@ -116,6 +153,7 @@ macro_rules! address { storage: PhantomData, } } + /// Return the raw address without extending its allocation lifetime. #[inline(always)] pub const fn as_raw(self) -> *mut T { self.pointer @@ -132,17 +170,24 @@ address!( "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> { @@ -213,6 +258,7 @@ impl<'s, F: PhysicalFormat> SharedRegion<'s, F> { 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/cuteir/ctm-rs/src/mma.rs b/cuda-oxide/cuteir/ctm-rs/src/mma.rs index c235c184d2..20d71536f3 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/mma.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/mma.rs @@ -2,7 +2,7 @@ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ -//! SM120 warp-level MXFP4 instructions and role-specific register carriers. +//! Warp-level FP16 and SM120 MXFP4 instructions with role-specific registers. //! Addresses, lane mapping, synchronization and issue order remain explicit. use super::memory::SharedAddress; @@ -25,6 +25,7 @@ use super::memory::SharedAddress; 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 })* @@ -32,17 +33,23 @@ macro_rules! static_for { } 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([$element; $count]); + 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 @@ -50,6 +57,18 @@ macro_rules! registers { } }; } +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!( Mxf4A, u32, @@ -69,12 +88,29 @@ registers!( "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) }) +} + /// One lane's rectangular collection of warp-wide 16x8 accumulator fragments. /// /// `ROWS` and `COLS` count fragments, not scalar elements. Physical placement is @@ -83,6 +119,7 @@ impl Accumulator16x8 { /// 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], } @@ -148,18 +185,20 @@ impl AccumulatorTile16x8<2, COLS> { } } +/// 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(u16); + 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 @@ -233,7 +272,7 @@ pub unsafe fn mma_mxf4( /// 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 = cuda_device::convert::cvt_f16x2_f32(c.0[0], c.0[1]); - let bottom = cuda_device::convert::cvt_f16x2_f32(c.0[2], c.0[3]); + 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/cuteir/ctm-rs/src/tcgen05/descriptors.rs b/cuda-oxide/cuteir/ctm-rs/src/tcgen05/descriptors.rs index edfafa453a..eccd11aa93 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/tcgen05/descriptors.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/tcgen05/descriptors.rs @@ -13,33 +13,64 @@ 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. pub enum KMajor {} +/// Left matrix input; cannot be substituted for [`OperandB`]. pub enum OperandA {} +/// Right matrix input; cannot be substituted for [`OperandA`]. 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 {} @@ -56,11 +87,15 @@ macro_rules! spec { 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 { @@ -68,6 +103,7 @@ impl InstructionDescriptor { spec: PhantomData, } } + /// Return the instruction’s packed 32-bit descriptor. #[inline(always)] pub const fn raw(&self) -> u32 { self.raw.raw() @@ -79,12 +115,16 @@ impl Default for InstructionDescriptor { } } +/// 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); @@ -99,18 +139,26 @@ fn encode_smem_offset(offset: ByteOffset) -> Result: 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 { @@ -146,18 +194,23 @@ impl<'s, R: OperandRole, F: MmaInputFormat> SmemOperand<'s, R, F> { } } #[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); @@ -183,6 +236,7 @@ mod tests { 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(); diff --git a/cuda-oxide/cuteir/ctm-rs/src/tcgen05/mma.rs b/cuda-oxide/cuteir/ctm-rs/src/tcgen05/mma.rs index fed7687aa9..3fde634a6d 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/tcgen05/mma.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/tcgen05/mma.rs @@ -10,17 +10,26 @@ use super::{ 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) { @@ -71,7 +80,10 @@ pub unsafe fn mma( /// Registers returned by a pending TMEM load; decode only after an explicit wait. #[repr(transparent)] -pub struct TmemRaw32(CuSimd); +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. /// @@ -99,6 +111,7 @@ impl TmemRaw32 { 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)),*]) diff --git a/cuda-oxide/cuteir/ctm-rs/src/tcgen05/mod.rs b/cuda-oxide/cuteir/ctm-rs/src/tcgen05/mod.rs index 2a0a68bb2d..7ee782ecb2 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/tcgen05/mod.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/tcgen05/mod.rs @@ -2,6 +2,13 @@ * 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. +//! +//! Start with [`descriptors::DenseSharedF16`] for supported shapes and +//! [`tmem::TmemGuard`] for allocation ownership. Issue, commit, wait and release +//! remain explicit. The crate's `docs/primitives-coverage.md` gives a complete +//! issue-site example and the supported compiler profiles. + pub mod descriptors; pub mod mma; pub mod tmem; diff --git a/cuda-oxide/cuteir/ctm-rs/src/tcgen05/tmem.rs b/cuda-oxide/cuteir/ctm-rs/src/tcgen05/tmem.rs index 0d047098ee..090ebaaa1f 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/tcgen05/tmem.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/tcgen05/tmem.rs @@ -20,7 +20,9 @@ pub use cuda_device::tcgen05::{TmemAddress, TmemDeallocated, TmemReady, TmemUnin /// 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)>, } @@ -28,29 +30,46 @@ pub struct TmemGuard { // 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(u32); +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> {} @@ -69,14 +88,19 @@ impl TmemFormat for TmemReadF32<32, 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, @@ -104,12 +128,18 @@ fn pack_address( /// 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. /// @@ -206,6 +236,8 @@ where _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, @@ -217,6 +249,18 @@ where 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, @@ -231,9 +275,11 @@ where } } +/// 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!( diff --git a/cuda-oxide/cuteir/ctm-rs/src/tma.rs b/cuda-oxide/cuteir/ctm-rs/src/tma.rs index 32c41bc20e..65cef45a41 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/tma.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/tma.rs @@ -16,16 +16,21 @@ 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 } @@ -126,16 +131,22 @@ pub unsafe fn prefetch>(map: &TensorMap) pub unsafe fn prefetch_descriptor(descriptor: *const TmaDescriptor) { unsafe { cuda_device::tma::prefetch_tma_descriptor(descriptor) } } +/// Commit this thread’s pending bulk copies as one group; does not wait. #[inline(always)] pub fn bulk_commit_group() { cuda_device::tma::cp_async_bulk_commit_group() } +/// Limits bulk-copy wait counts to the instruction’s immediate range. 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(); } } +/// 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. @@ -154,7 +165,9 @@ use core::marker::PhantomData; /// 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 {} @@ -177,6 +190,7 @@ impl> TensorMap { 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/cuteir/ctm-rs/src/warp.rs b/cuda-oxide/cuteir/ctm-rs/src/warp.rs index 546555cc81..5a0179dc1f 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/warp.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/warp.rs @@ -2,7 +2,18 @@ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ -//! Explicit lane identity and election. No other primitive elects implicitly. +//! 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+). diff --git a/cuda-oxide/cuteir/ctm-rs/tests/accumulator_tile.rs b/cuda-oxide/cuteir/ctm-rs/tests/accumulator_tile.rs index 956e153301..003f94cd4d 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/accumulator_tile.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/accumulator_tile.rs @@ -1,7 +1,9 @@ // 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(); @@ -26,8 +28,10 @@ fn static_fragment_access_preserves_bands_columns_and_register_order() { ); } +/// 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) } @@ -40,6 +44,7 @@ fn static_expansion_nests_in_source_order_with_const_generic_indices() { 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(); diff --git a/cuda-oxide/cuteir/ctm-rs/tests/attention_api.rs b/cuda-oxide/cuteir/ctm-rs/tests/attention_api.rs new file mode 100644 index 0000000000..08e3acbedd --- /dev/null +++ b/cuda-oxide/cuteir/ctm-rs/tests/attention_api.rs @@ -0,0 +1,11 @@ +/* + * 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() { + trybuild::TestCases::new().pass("tests/compile_pass/attention_primitives.rs"); +} diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_bounds.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_bounds.rs index e20b1d5ddb..fe30f6f5cf 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_bounds.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_bounds.rs @@ -1,10 +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/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.rs index 3ef913d724..0cd3f9110d 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.rs @@ -1,7 +1,10 @@ // 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 { diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.stderr index 106b9e8478..e568182aec 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.stderr @@ -1,23 +1,23 @@ error[E0308]: mismatched types - --> tests/compile_fail/primitives_accumulator_tile_roles.rs:9:14 - | -9 | [Mxf4B::from_registers([0; 2]); 2], - | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `Mxf4A`, found `Mxf4B` + --> 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:10:14 + --> tests/compile_fail/primitives_accumulator_tile_roles.rs:13:14 | -10 | [Mxf4A::from_registers([0; 4]); 2], +13 | [Mxf4A::from_registers([0; 4]); 2], | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `Mxf4B`, found `Mxf4A` error[E0308]: mismatched types - --> tests/compile_fail/primitives_accumulator_tile_roles.rs:11:14 + --> tests/compile_fail/primitives_accumulator_tile_roles.rs:14:14 | -11 | [ScalePairB::from_bits(0); 2], +14 | [ScalePairB::from_bits(0); 2], | ^^^^^^^^^^^^^^^^^^^^^^^^ expected `ScalePairA`, found `ScalePairB` error[E0308]: mismatched types - --> tests/compile_fail/primitives_accumulator_tile_roles.rs:12:14 + --> tests/compile_fail/primitives_accumulator_tile_roles.rs:15:14 | -12 | [ScalePairA::from_bits(0); 2], +15 | [ScalePairA::from_bits(0); 2], | ^^^^^^^^^^^^^^^^^^^^^^^^ expected `ScalePairB`, found `ScalePairA` diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.rs index cb4b95364f..00781c7bf9 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.rs @@ -1,7 +1,10 @@ // 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 { diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.stderr index f6c215de5f..131426e346 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.stderr @@ -1,7 +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:8:14 - | -8 | tile.mma_mxf4_pair::<0>( - | ^^^^^^^^^^^^^ method not found in `AccumulatorTile16x8<1, 8>` - | - = note: the method was found for `AccumulatorTile16x8<2, COLS>` + --> 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/cuteir/ctm-rs/tests/compile_fail/primitives_address_domain.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_address_domain.rs index 6118900c00..2fc24b349f 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_address_domain.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_address_domain.rs @@ -2,7 +2,11 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_address_domain.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_address_domain.stderr index c213d9d960..e29a79ca9e 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_address_domain.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_address_domain.stderr @@ -1,15 +1,15 @@ error[E0308]: mismatched types - --> tests/compile_fail/primitives_address_domain.rs:7:52 - | -7 | 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>` + --> 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:6:4 - | -6 | fn cluster(_: ClusterSharedAddress<'_,u32>) {} - | ^^^^^^^ ------------------------------- + --> tests/compile_fail/primitives_address_domain.rs:8:4 + | + 8 | fn cluster(_: ClusterSharedAddress<'_,u32>) {} + | ^^^^^^^ ------------------------------- diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_async_wait_runtime.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_async_wait_runtime.rs new file mode 100644 index 0000000000..1944cfbf2e --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/ctm-rs/tests/compile_fail/primitives_async_wait_runtime.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_async_wait_runtime.stderr new file mode 100644 index 0000000000..323273f924 --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.rs index 92f92f3c3f..e0c7144899 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.rs @@ -2,6 +2,9 @@ * 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::bulk_wait_group_read; +/// Reject a bulk-copy wait count above seven. fn wrong() { 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/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.stderr index 1291de7fda..ff1424b001 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.stderr @@ -1,7 +1,7 @@ error[E0277]: the trait bound `(): ValidWaitCount<8>` is not satisfied - --> tests/compile_fail/primitives_bulk_wait_count.rs:6:37 + --> tests/compile_fail/primitives_bulk_wait_count.rs:8:37 | -6 | fn wrong() { bulk_wait_group_read::<8>(); } +8 | fn wrong() { bulk_wait_group_read::<8>(); } | ^ the trait `ValidWaitCount<8>` is not implemented for `()` | = help: `()` implements trait `ValidWaitCount`: diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.rs index e05d2873fb..f29e82e67c 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.rs @@ -2,6 +2,9 @@ * 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::bulk_wait_group_read; +/// Require bulk-copy wait counts as const generics rather than runtime arguments. fn wrong(n:u32) { 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/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.stderr index dd35225d12..a82c1157b9 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.stderr @@ -1,7 +1,7 @@ error[E0061]: this function takes 0 arguments but 1 argument was supplied - --> tests/compile_fail/primitives_bulk_wait_runtime.rs:6:19 + --> tests/compile_fail/primitives_bulk_wait_runtime.rs:8:19 | -6 | fn wrong(n:u32) { bulk_wait_group_read(n); } +8 | fn wrong(n:u32) { bulk_wait_group_read(n); } | ^^^^^^^^^^^^^^^^^^^^ - unexpected argument of type `u32` | note: function defined here @@ -11,6 +11,6 @@ note: function defined here | ^^^^^^^^^^^^^^^^^^^^ help: remove the extra argument | -6 - fn wrong(n:u32) { bulk_wait_group_read(n); } -6 + fn wrong(n:u32) { bulk_wait_group_read(); } +8 - fn wrong(n:u32) { bulk_wait_group_read(n); } +8 + fn wrong(n:u32) { bulk_wait_group_read(); } | diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_f16_mma_roles.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_f16_mma_roles.rs new file mode 100644 index 0000000000..8ae2553fc0 --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/ctm-rs/tests/compile_fail/primitives_f16_mma_roles.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_f16_mma_roles.stderr new file mode 100644 index 0000000000..1b3fabc7bd --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/ctm-rs/tests/compile_fail/primitives_guard_consumed.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_guard_consumed.rs index 2ae82aa713..0e3302dae9 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_guard_consumed.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_guard_consumed.rs @@ -2,6 +2,9 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_guard_consumed.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_guard_consumed.stderr index 510aac71c3..95999c7b85 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_guard_consumed.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_guard_consumed.stderr @@ -1,7 +1,7 @@ error[E0382]: borrow of moved value: `g` - --> tests/compile_fail/primitives_guard_consumed.rs:6:101 + --> tests/compile_fail/primitives_guard_consumed.rs:8:101 | -6 | fn wrong(g:TmemGuard) { let _ = unsafe { g.dealloc_collective() }; let _ = g.raw_address(); } +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 diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_invalid_columns.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_invalid_columns.rs index 0bf1cf5fd6..42f77f0bf2 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_invalid_columns.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_invalid_columns.rs @@ -2,6 +2,9 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_invalid_columns.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_invalid_columns.stderr index 5a45e6d36c..d78b042085 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_invalid_columns.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_invalid_columns.stderr @@ -1,7 +1,7 @@ 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:6:82 + --> tests/compile_fail/primitives_invalid_columns.rs:8:82 | -6 | fn wrong(slot:*mut u32) { let _ = unsafe { TmemGuard::::from_raw_slot(slot) }; } +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: diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_b_shape.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_b_shape.rs index b6ed681aa7..fdbb520f0d 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_b_shape.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_b_shape.rs @@ -2,11 +2,16 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_mma_b_shape.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_b_shape.stderr index 1022f25fac..d6b7774669 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_b_shape.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_b_shape.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/primitives_mma_b_shape.rs:10:45 + --> tests/compile_fail/primitives_mma_b_shape.rs:14:45 | -10 | unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; +14 | unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; | ----------------------------- ^ expected `96`, found `80` | | | arguments to this function are incorrect diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_configuration.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_configuration.rs index ad0bccf165..26ae19312a 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_configuration.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_configuration.rs @@ -2,14 +2,20 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_mma_configuration.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_configuration.stderr index bb3860ca76..a328f98b84 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_configuration.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_configuration.stderr @@ -1,7 +1,7 @@ error[E0277]: the trait bound `ctm_rs::tcgen05::descriptors::DenseSharedF16: MmaSpec` is not satisfied - --> tests/compile_fail/primitives_mma_configuration.rs:11:10 + --> tests/compile_fail/primitives_mma_configuration.rs:16:10 | -11 | let _ = InstructionDescriptor::>::new(); +16 | let _ = InstructionDescriptor::>::new(); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ unsatisfied trait bound | = help: the trait `MmaSpec` is not implemented for `ctm_rs::tcgen05::descriptors::DenseSharedF16` @@ -26,9 +26,9 @@ note: required by a bound in `ctm_rs::tcgen05::descriptors::InstructionDescripto = 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:11:84 + --> tests/compile_fail/primitives_mma_configuration.rs:16:84 | -11 | let _ = InstructionDescriptor::>::new(); +16 | let _ = InstructionDescriptor::>::new(); | ^^^ associated function or constant cannot be called due to unsatisfied trait bounds | ::: src/tcgen05/descriptors.rs @@ -40,9 +40,9 @@ error[E0599]: the associated function or constant `new` exists for struct `ctm_r `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:12:10 + --> tests/compile_fail/primitives_mma_configuration.rs:17:10 | -12 | let _ = InstructionDescriptor::>::new(); +17 | let _ = InstructionDescriptor::>::new(); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ unsatisfied trait bound | = help: the trait `MmaSpec` is not implemented for `ctm_rs::tcgen05::descriptors::DenseSharedF16` @@ -67,9 +67,9 @@ note: required by a bound in `ctm_rs::tcgen05::descriptors::InstructionDescripto = 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:12:84 + --> tests/compile_fail/primitives_mma_configuration.rs:17:84 | -12 | let _ = InstructionDescriptor::>::new(); +17 | let _ = InstructionDescriptor::>::new(); | ^^^ associated function or constant cannot be called due to unsatisfied trait bounds | ::: src/tcgen05/descriptors.rs @@ -81,9 +81,9 @@ error[E0599]: the associated function or constant `new` exists for struct `ctm_r `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:13:10 + --> tests/compile_fail/primitives_mma_configuration.rs:18:10 | -13 | let _ = InstructionDescriptor::>::new(); +18 | let _ = InstructionDescriptor::>::new(); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ unsatisfied trait bound | = help: the trait `MmaSpec` is not implemented for `ctm_rs::tcgen05::descriptors::DenseSharedF16` @@ -108,9 +108,9 @@ note: required by a bound in `ctm_rs::tcgen05::descriptors::InstructionDescripto = 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:13:85 + --> tests/compile_fail/primitives_mma_configuration.rs:18:85 | -13 | let _ = InstructionDescriptor::>::new(); +18 | let _ = InstructionDescriptor::>::new(); | ^^^ associated function or constant cannot be called due to unsatisfied trait bounds | ::: src/tcgen05/descriptors.rs diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_descriptor.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_descriptor.rs index 8cbfbd051c..f66ae38120 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_descriptor.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_descriptor.rs @@ -2,11 +2,16 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_mma_descriptor.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_descriptor.stderr index f1a27efd95..e98583d409 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_descriptor.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_descriptor.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/primitives_mma_descriptor.rs:10:47 + --> tests/compile_fail/primitives_mma_descriptor.rs:14:47 | -10 | unsafe { mma::(d,a,b,&InstructionDescriptor::::new(),false) }; +14 | unsafe { mma::(d,a,b,&InstructionDescriptor::::new(),false) }; | ----------------------------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `192`, found `160` | | | arguments to this function are incorrect diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_destination.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_destination.rs index 13172267c1..2ce590baa0 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_destination.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_destination.rs @@ -2,11 +2,16 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_mma_destination.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_destination.stderr index 561732fd77..6e600320ac 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_destination.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_destination.stderr @@ -1,7 +1,7 @@ error[E0271]: type mismatch resolving `> as TmemOperand>::Format == TmemF32<256, 192>` - --> tests/compile_fail/primitives_mma_destination.rs:10:41 + --> tests/compile_fail/primitives_mma_destination.rs:14:41 | -10 | unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; +14 | unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; | ----------------------------- ^ expected `192`, found `160` | | | required by a bound introduced by this call diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_k_step.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_k_step.rs index c83a03579f..f2d42eac7c 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_k_step.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_k_step.rs @@ -2,9 +2,14 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_mma_k_step.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_k_step.stderr index b1607d439c..9312efd2dc 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_k_step.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_k_step.stderr @@ -1,28 +1,28 @@ error[E0277]: the trait bound `(): ValidKStep<4>` is not satisfied - --> tests/compile_fail/primitives_mma_k_step.rs:9:65 - | -9 | fn wrong(a:&SmemOperand) { let _ = a.k16::<4>(); } - | ^ the trait `ValidKStep<4>` is not implemented for `()` - | + --> 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 + --> 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 - | - | impl<'s, R: OperandRole, F: MmaInputFormat> SmemOperand<'s, R, F> { + --> src/tcgen05/descriptors.rs + | + | impl<'s, R: OperandRole, F: MmaInputFormat> SmemOperand<'s, R, F> { ... - | pub fn k16(&self) -> Self - | 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) + | pub fn k16(&self) -> Self + | 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/cuteir/ctm-rs/tests/compile_fail/primitives_mma_output_format.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_output_format.rs index 080cf2011e..ea12f5bf92 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_output_format.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_output_format.rs @@ -2,9 +2,14 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_mma_output_format.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_output_format.stderr index 68fbcf5716..134227ce3a 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_output_format.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_output_format.stderr @@ -1,37 +1,37 @@ error[E0277]: the trait bound `Swizzled2D<128, 32, Swizzle64>: MmaInputFormat` is not satisfied - --> tests/compile_fail/primitives_mma_output_format.rs:9:47 - | -9 | fn wrong(a:&SharedRegion) { let _ = SmemOperand::::from_region(a); } - | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the trait `MmaInputFormat` is not implemented for `Swizzled2D<128, 32, Swizzle64>` - | + --> 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 + --> 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) + --> 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:9:81 - | -9 | 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(PhantomData); - | -------------------------------------------------------------------- doesn't satisfy `Swizzled2D<128, 32, Swizzle64>: MmaInputFormat` - | - = note: the following trait bounds were not satisfied: - `Swizzled2D<128, 32, Swizzle64>: MmaInputFormat` + --> 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/cuteir/ctm-rs/tests/compile_fail/primitives_mma_roles.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_roles.rs index 7bd91227c5..8bbac278fe 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_roles.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_roles.rs @@ -2,11 +2,16 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_mma_roles.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_roles.stderr index 5ffe4487c1..63f0bfe4ad 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_roles.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_mma_roles.stderr @@ -1,20 +1,20 @@ error[E0308]: arguments to this function are incorrect - --> tests/compile_fail/primitives_mma_roles.rs:10:11 + --> tests/compile_fail/primitives_mma_roles.rs:14:11 | -10 | unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; +14 | unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ | note: expected `&SmemOperand<'_, OperandA, _>`, found `&SmemOperand<'_, OperandB, _>` - --> tests/compile_fail/primitives_mma_roles.rs:10:43 + --> tests/compile_fail/primitives_mma_roles.rs:14:43 | -10 | unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; +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:10:45 + --> tests/compile_fail/primitives_mma_roles.rs:14:45 | -10 | unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; +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>>` diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.rs index d04b70fab7..34f4d8d0ca 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.rs @@ -2,7 +2,11 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.stderr index cf4a704725..dc48b33275 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.stderr @@ -1,11 +1,11 @@ error[E0599]: no method named `alloc` found for struct `TmemGuard` in the current scope - --> tests/compile_fail/primitives_no_implicit_allocation.rs:6:71 + --> tests/compile_fail/primitives_no_implicit_allocation.rs:8:71 | -6 | fn group1(g:TmemGuard) { let _ = unsafe { g.alloc() }; } +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:7:71 - | -7 | fn group2(g:TmemGuard) { let _ = unsafe { g.alloc() }; } - | ^^^^^ method not found in `TmemGuard` + --> 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/cuteir/ctm-rs/tests/compile_fail/primitives_remote_domain.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_domain.rs index 9f310d8806..126b54f139 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_domain.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_domain.rs @@ -2,7 +2,10 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_remote_domain.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_domain.stderr index aafc72679d..defc28ec88 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_domain.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_domain.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/primitives_remote_domain.rs:7:65 + --> tests/compile_fail/primitives_remote_domain.rs:9:65 | -7 | fn wrong(local:SharedAddress) { unsafe { arrive_remote(local) }; } +9 | fn wrong(local:SharedAddress) { unsafe { arrive_remote(local) }; } | ------------- ^^^^^ expected `ClusterSharedAddress<'_, Barrier>`, found `SharedAddress<'_, Barrier>` | | | arguments to this function are incorrect diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.rs index 5b09b88092..6329b3d9fb 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.rs @@ -2,9 +2,12 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.stderr index 31ecb1b721..abfc2f8a92 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.stderr @@ -1,15 +1,15 @@ error[E0308]: mismatched types - --> tests/compile_fail/primitives_remote_relaxed_domain.rs:8:36 - | -8 | 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>` + --> 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>) { - | ^^^^^^^^^^^^^^^^^^^^^ + --> src/cluster.rs + | + | pub unsafe fn arrive_remote_relaxed(address: ClusterSharedAddress<'_, Barrier>) { + | ^^^^^^^^^^^^^^^^^^^^^ diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tensor_map_element.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tensor_map_element.rs index 0e264dca1b..2af3e7c34f 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tensor_map_element.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tensor_map_element.rs @@ -2,7 +2,10 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_tensor_map_element.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tensor_map_element.stderr index 5fd75c77f0..50f6c8b8c4 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tensor_map_element.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tensor_map_element.stderr @@ -1,7 +1,7 @@ error[E0271]: type mismatch resolving ` as PhysicalFormat>::Element == f32` - --> tests/compile_fail/primitives_tensor_map_element.rs:7:13 + --> tests/compile_fail/primitives_tensor_map_element.rs:9:13 | -7 | fn wrong(_: &TensorMap) {} +9 | fn wrong(_: &TensorMap) {} | ^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `f32`, found `f16` | note: required by a bound in `TensorMap` diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tma_format.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tma_format.rs index fa69e07938..b10517cfe9 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tma_format.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tma_format.rs @@ -2,7 +2,10 @@ * 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 { store_2d(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/cuteir/ctm-rs/tests/compile_fail/primitives_tma_format.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tma_format.stderr index 3bca5a75fa..3ff1833999 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tma_format.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tma_format.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/primitives_tma_format.rs:7:94 + --> tests/compile_fail/primitives_tma_format.rs:9:94 | -7 | fn wrong(src:&SharedRegion,map:&TensorMap) { unsafe { store_2d(src,map,Coordinates2D::new(0,0)) }; } +9 | fn wrong(src:&SharedRegion,map:&TensorMap) { unsafe { store_2d(src,map,Coordinates2D::new(0,0)) }; } | -------- ^^^ expected `32`, found `64` | | | arguments to this function are incorrect diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_unready_barrier.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_unready_barrier.rs index 2ec543288d..1d84b47b26 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_unready_barrier.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_unready_barrier.rs @@ -2,7 +2,10 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_unready_barrier.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_unready_barrier.stderr index 41ab39e9b8..e6077af33f 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_unready_barrier.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_unready_barrier.stderr @@ -1,5 +1,5 @@ error[E0599]: no method named `shared_address` found for struct `ManagedBarrier` in the current scope - --> tests/compile_fail/primitives_unready_barrier.rs:7:59 + --> tests/compile_fail/primitives_unready_barrier.rs:9:59 | -7 | fn wrong(b:ManagedBarrier) { let _ = b.shared_address(); } +9 | fn wrong(b:ManagedBarrier) { let _ = b.shared_address(); } | ^^^^^^^^^^^^^^ method not found in `ManagedBarrier` diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.rs index c0ef7b2fb8..171b511e97 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.rs @@ -2,11 +2,15 @@ * 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/cuteir/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.stderr index 617f4522b2..1be02f4a62 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.stderr +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.stderr @@ -1,11 +1,11 @@ error[E0505]: cannot move out of `g` because it is borrowed - --> tests/compile_fail/primitives_view_after_deallocation.rs:9:26 + --> tests/compile_fail/primitives_view_after_deallocation.rs:12:26 | - 7 | fn wrong(g:TmemGuard) { +10 | fn wrong(g:TmemGuard) { | - binding `g` declared here - 8 | let view = g.view::>(ColumnOffset::new(0)).unwrap(); +11 | let view = g.view::>(ColumnOffset::new(0)).unwrap(); | - borrow of `g` occurs here - 9 | let _dead = unsafe { g.dealloc_collective() }; +12 | let _dead = unsafe { g.dealloc_collective() }; | ^ move out of `g` occurs here -10 | consume_view(view); +13 | consume_view(view); | ---- borrow later used here diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_pass/attention_primitives.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_pass/attention_primitives.rs new file mode 100644 index 0000000000..599f608bcf --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/ctm-rs/tests/primitives_mxf4_storage.rs b/cuda-oxide/cuteir/ctm-rs/tests/primitives_mxf4_storage.rs index f43629c5c8..2cbf9e159a 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/primitives_mxf4_storage.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/primitives_mxf4_storage.rs @@ -2,6 +2,7 @@ * 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, @@ -11,6 +12,7 @@ use ctm_rs::{ }; 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); @@ -26,6 +28,7 @@ fn blockscaled_stage_storage_matches_instruction_footprints() { 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 { @@ -67,6 +70,7 @@ fn blockscaled_storage_rejects_short_or_misaligned_stages() { } } +/// 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)] @@ -85,11 +89,16 @@ fn ordinary_tma_load_uses_cta_local_completion_for_both_stage_types() { ) = load_2d::; } +/// 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([u8; 32768]); + 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, diff --git a/cuda-oxide/cuteir/ctm-rs/tests/primitives_safety.rs b/cuda-oxide/cuteir/ctm-rs/tests/primitives_safety.rs index c307e7fcf4..8a461409a1 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/primitives_safety.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/primitives_safety.rs @@ -2,6 +2,8 @@ * 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/cuteir/ctm-rs/tests/primitives_sm120.rs b/cuda-oxide/cuteir/ctm-rs/tests/primitives_sm120.rs index 2ac42da629..ae76090995 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/primitives_sm120.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/primitives_sm120.rs @@ -1,6 +1,8 @@ // 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!( diff --git a/cuda-oxide/cuteir/ctm-rs/tests/primitives_types.rs b/cuda-oxide/cuteir/ctm-rs/tests/primitives_types.rs index 6cafc9a8c4..87bee7c6bf 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/primitives_types.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/primitives_types.rs @@ -2,9 +2,11 @@ * 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); @@ -15,10 +17,12 @@ fn physical_formats_and_tensor_map_preserve_kernel_abi() { 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. @@ -46,6 +50,7 @@ fn shared_region_rejects_invalid_storage_before_use() { ); } +/// 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}; @@ -58,6 +63,7 @@ fn multicast_masks_and_signed_coordinates_preserve_instruction_domains() { 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; diff --git a/cuda-oxide/cuteir/cute-layout/src/algebra.rs b/cuda-oxide/cuteir/cute-layout/src/algebra.rs index b557b3ba61..52033e61b2 100644 --- a/cuda-oxide/cuteir/cute-layout/src/algebra.rs +++ b/cuda-oxide/cuteir/cute-layout/src/algebra.rs @@ -80,7 +80,9 @@ use crate::{IntTuple, Layout}; /// 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(pub String); +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 { @@ -90,6 +92,7 @@ impl fmt::Display for AlgebraError { 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. @@ -320,6 +323,7 @@ fn compose_int(a: &Layout, b_shape: i64, b_stride: i64) -> Result { 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 } @@ -959,6 +963,7 @@ pub fn max_alignment(layout: &Layout) -> Result { 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(); @@ -970,11 +975,13 @@ fn gcd_nonnegative(a: i64, b: i64) -> i64 { a as i64 } +/// 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() } @@ -1000,6 +1007,7 @@ mod tests { .collect() } + /// 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"); @@ -1067,6 +1075,7 @@ mod tests { ); } + /// 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 @@ -1216,6 +1225,7 @@ mod tests { assert_eq!(seen, expect); } + /// 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 @@ -1226,6 +1236,7 @@ mod tests { assert_eq!(seen, (0..24).collect::>()); } + /// Static alignment factors must agree with the pinned cutegen examples. #[test] fn max_alignment_matches_cutegen_static_examples() { let cases = [ @@ -1241,6 +1252,7 @@ mod tests { } } + /// 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/cuteir/cute-layout/src/composed.rs b/cuda-oxide/cuteir/cute-layout/src/composed.rs index b0e95c0c78..7d9d6443e6 100644 --- a/cuda-oxide/cuteir/cute-layout/src/composed.rs +++ b/cuda-oxide/cuteir/cute-layout/src/composed.rs @@ -111,7 +111,7 @@ pub enum ComposedLayoutError { InvalidSwizzle, /// Element-to-byte conversion needs a positive power-of-two byte size. /// The payload is the rejected size. - InvalidElementBytes(i64), + InvalidElementBytes(#[doc = "Rejected byte size for one element."] i64), /// Scaling the offset or a stride exceeded the `i64` offset model. ArithmeticOverflow, } @@ -143,9 +143,13 @@ impl core::error::Error for ComposedLayoutError {} /// the allocation's base pointer. #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct ComposedLayout { + /// Offset-bit permutation applied after the inner map and fixed displacement. outer: Swizzle, + /// Fixed displacement in the recorded unit, added before the swizzle. offset: i64, + /// Coordinate-to-offset map evaluated before the outer swizzle. inner: Layout, + /// Whether offsets and strides count elements or bytes. unit: OffsetUnit, } @@ -378,6 +382,7 @@ impl fmt::Display for ComposedLayout { } } +/// 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), @@ -389,6 +394,7 @@ fn checked_scale_tuple(tuple: &IntTuple, scale: i64) -> Option { } } +/// 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(); @@ -398,6 +404,7 @@ fn gcd(a: i64, b: i64) -> i64 { 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 @@ -405,15 +412,18 @@ fn unsigned_to_i64(value: u64) -> i64 { 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 inner layout, fixed displacement and outer swizzle in that order. #[test] fn call_uses_cutegen_composition_order() { let layout = @@ -429,6 +439,7 @@ mod tests { 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( @@ -450,6 +461,7 @@ mod tests { } } + /// 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); @@ -466,6 +478,7 @@ mod tests { ); } + /// 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 @@ -482,6 +495,7 @@ mod tests { ); } + /// Checked composed evaluation must return failure when the fixed displacement overflows. #[test] fn checked_call_rejects_offset_overflow() { let layout = @@ -489,6 +503,7 @@ mod tests { assert_eq!(layout.checked_call(&IntTuple::Leaf(1)), None); } + /// Alignment must account for inner strides, fixed displacement and swizzle-protected bits. #[test] fn max_alignment_combines_every_part() { let layout = ComposedLayout::new( @@ -519,6 +534,7 @@ mod tests { assert_eq!(inner_limited.max_alignment().unwrap(), 4); } + /// An identity swizzle must not reduce alignment already guaranteed by the inner 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 diff --git a/cuda-oxide/cuteir/cute-layout/src/lib.rs b/cuda-oxide/cuteir/cute-layout/src/lib.rs index 31611ab498..29c28769ff 100644 --- a/cuda-oxide/cuteir/cute-layout/src/lib.rs +++ b/cuda-oxide/cuteir/cute-layout/src/lib.rs @@ -5,6 +5,16 @@ //! 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 @@ -199,9 +209,9 @@ use core::str::FromStr; #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum IntTuple { /// A single integer with no nested structure. - Leaf(i64), + Leaf(#[doc = "One shape, stride or coordinate integer."] i64), /// An ordered group of integers or further groups. - Tuple(Vec), + Tuple(#[doc = "Ordered child modes whose nested grouping is preserved."] Vec), } impl IntTuple { @@ -254,6 +264,7 @@ impl IntTuple { 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), @@ -395,6 +406,7 @@ fn checked_crd2idx(coord: &IntTuple, shape: &IntTuple, stride: &IntTuple) -> Opt } } +/// 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), @@ -745,6 +757,7 @@ impl fmt::Display for ParseLayoutError { 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('(') { @@ -801,15 +814,18 @@ impl FromStr for Layout { } } +/// Regression tests for this module’s layout, representation and rejection contracts. #[cfg(test)] mod tests { 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 @@ -821,6 +837,7 @@ mod tests { 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 @@ -831,6 +848,7 @@ mod tests { 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; @@ -842,6 +860,7 @@ mod tests { 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()); @@ -862,6 +881,7 @@ mod tests { assert_eq!(overflowing_address.checked_call(&IntTuple::Leaf(2)), None); } + /// 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 diff --git a/cuda-oxide/cuteir/cute-layout/src/predication.rs b/cuda-oxide/cuteir/cute-layout/src/predication.rs index a9d9b8c2df..46f72851c0 100644 --- a/cuda-oxide/cuteir/cute-layout/src/predication.rs +++ b/cuda-oxide/cuteir/cute-layout/src/predication.rs @@ -66,23 +66,11 @@ //! zero read: base + 0 (aligned fallback; source_bytes = 0) //! ``` //! -//! This module computes data only: the compile-time half of the split. The -//! predicate's value depends on runtime numbers, so the emitted kernel -//! evaluates it; what is settled here is which comparisons exist and why -//! every branch-free outcome is memory-safe: -//! -//! ```text -//! this module (compile time) emitted kernel (GPU runtime) -//! ────────────────────────── ──────────────────────────── -//! which comparisons must exist evaluate them: row < rows? -//! prove base + 0 is a safe fallback pick the offset with select -//! reject pitch < cols, overflow, issue cp.async with the chosen -//! bad atom widths, loudly source_bytes (16 / 8 / 0) -//! ``` -//! -//! Later, the compiler must use the plan to choose the offset and source -//! size consistently for every participating thread, and then issue -//! `cp.async` for each one. +//! 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; @@ -156,7 +144,7 @@ pub type PredicationResult = core::result::Result; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PredicationError { /// `cp.async` supports only 4-, 8-, or 16-byte transaction widths here. - UnsupportedAtomBytes(u32), + 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. @@ -344,10 +332,12 @@ pub fn plan_uniform_cp_async_source( }) } +/// 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, @@ -355,6 +345,7 @@ mod tests { 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!( @@ -366,6 +357,7 @@ mod tests { ); } + /// 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!( @@ -377,6 +369,7 @@ mod tests { ); } + /// 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 [ @@ -396,6 +389,7 @@ mod tests { } } + /// In-bounds coordinates must still reject element-offset and byte-offset overflow. #[test] fn valid_offset_arithmetic_is_checked() { let huge_pitch = GmemMatrixBounds { @@ -424,6 +418,7 @@ mod tests { ); } + /// 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!( @@ -455,6 +450,7 @@ mod tests { ); } + /// 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 { diff --git a/cuda-oxide/cuteir/cute-layout/src/swizzle.rs b/cuda-oxide/cuteir/cute-layout/src/swizzle.rs index 690ba13065..962176a737 100644 --- a/cuda-oxide/cuteir/cute-layout/src/swizzle.rs +++ b/cuda-oxide/cuteir/cute-layout/src/swizzle.rs @@ -284,10 +284,12 @@ impl fmt::Display for Swizzle { } } +/// 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 @@ -305,6 +307,7 @@ mod tests { 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)); @@ -318,12 +321,14 @@ mod tests { 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 { @@ -331,6 +336,7 @@ mod tests { } } + /// 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 @@ -348,6 +354,7 @@ mod tests { 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); @@ -360,6 +367,7 @@ mod tests { } } + /// 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); @@ -368,6 +376,7 @@ mod tests { 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!( diff --git a/cuda-oxide/cuteir/cute-layout/src/validation.rs b/cuda-oxide/cuteir/cute-layout/src/validation.rs deleted file mode 100644 index 4bfad279e7..0000000000 --- a/cuda-oxide/cuteir/cute-layout/src/validation.rs +++ /dev/null @@ -1,1668 +0,0 @@ -/* - * 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. - -use alloc::collections::BTreeMap; -use alloc::format; -use alloc::string::String; -use alloc::vec; -use alloc::vec::Vec; -use core::fmt; - -use crate::algebra::make_layout_tv; -use crate::{ComposedLayout, IntTuple, Layout, OffsetUnit}; - -/// 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; - -/// 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(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 {} - -/// 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(&'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 {} - -fn plan_error(message: impl Into) -> CooperativePlanError { - CooperativePlanError(message.into()) -} - -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.leaves() { - 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(()) -} - -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_inner_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.inner(), "shared-memory inner")?; - 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 inner", smem_layout.inner()), - ] { - 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, - }) -} - -/// 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`] 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 -/// ``` -/// -/// `CpAsync` 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 `CpAsync`, 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 - .inner() - .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_inner_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(()) -} - -fn tv_coord(thread: i64, value: i64) -> IntTuple { - IntTuple::Tuple(vec![IntTuple::Leaf(thread), IntTuple::Leaf(value)]) -} - -fn positive_count(name: &'static str, value: i64) -> ValidationResult { - if value <= 0 { - Err(ValidationError::InvalidCount { name, value }) - } else { - Ok(value) - } -} - -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 }) -} - -/// 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(()) -} - -/// 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.inner().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.leaves(); - let col_leaves = modes[1].stride.leaves(); - 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.outer(); - if swizzle.bits == 0 { - return Ok(()); - } - let m_bytes = i64::from(swizzle.base) + element_bytes.trailing_zeros() as i64; - if swizzle.bits > 3 || m_bytes != 4 || swizzle.shift != 3 { - return Err(ValidationError::InvalidCount { - name: "swizzle is not one of TMA's encodable byte-unit S modes", - value: i64::from(swizzle.bits), - }); - } - let pitch_bytes = columns * element_bytes; - if pitch_bytes != (16 << swizzle.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.outer(); - if swizzle.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.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", - )) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::Swizzle; - use alloc::string::ToString; - - fn l(text: &str) -> Layout { - text.parse().unwrap() - } - - 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) - } - - #[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); - } - - #[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}"); - } - - #[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}"); - } - - #[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 { .. }) - )); - } - - #[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 { .. }) - )); - } - - #[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(); - } - - #[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", .. }) - )); - } - - #[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(); - } - - #[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(); - } - - #[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 { .. }) - )); - } - - #[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 { .. }) - )); - } - - #[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" - )) - ); - } - - #[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)); - } - - #[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)); - } - - #[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)); - } - - #[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/cuteir/cute-layout/src/validation/cooperative.rs b/cuda-oxide/cuteir/cute-layout/src/validation/cooperative.rs new file mode 100644 index 0000000000..669f2edf77 --- /dev/null +++ b/cuda-oxide/cuteir/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.leaves() { + 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_inner_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.inner(), "shared-memory inner")?; + 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 inner", smem_layout.inner()), + ] { + 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/cuteir/cute-layout/src/validation/error.rs b/cuda-oxide/cuteir/cute-layout/src/validation/error.rs new file mode 100644 index 0000000000..fcc41aed62 --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/cute-layout/src/validation/ldmatrix.rs b/cuda-oxide/cuteir/cute-layout/src/validation/ldmatrix.rs new file mode 100644 index 0000000000..db98a57a5a --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/cute-layout/src/validation/mod.rs b/cuda-oxide/cuteir/cute-layout/src/validation/mod.rs new file mode 100644 index 0000000000..72a9f53e76 --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/cute-layout/src/validation/tests.rs b/cuda-oxide/cuteir/cute-layout/src/validation/tests.rs new file mode 100644 index 0000000000..f9f1f66577 --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/cute-layout/src/validation/thread_value.rs b/cuda-oxide/cuteir/cute-layout/src/validation/thread_value.rs new file mode 100644 index 0000000000..68bcb56da2 --- /dev/null +++ b/cuda-oxide/cuteir/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 +/// ``` +/// +/// `CpAsync` 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 `CpAsync`, 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 + .inner() + .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_inner_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/cuteir/cute-layout/src/validation/tma.rs b/cuda-oxide/cuteir/cute-layout/src/validation/tma.rs new file mode 100644 index 0000000000..2636839522 --- /dev/null +++ b/cuda-oxide/cuteir/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.inner().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.leaves(); + let col_leaves = modes[1].stride.leaves(); + 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.outer(); + if swizzle.bits == 0 { + return Ok(()); + } + let m_bytes = i64::from(swizzle.base) + element_bytes.trailing_zeros() as i64; + if swizzle.bits > 3 || m_bytes != 4 || swizzle.shift != 3 { + return Err(ValidationError::InvalidCount { + name: "swizzle is not one of TMA's encodable byte-unit S modes", + value: i64::from(swizzle.bits), + }); + } + let pitch_bytes = columns * element_bytes; + if pitch_bytes != (16 << swizzle.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.outer(); + if swizzle.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.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/cuteir/cute-rs/examples/global_store_compile.rs b/cuda-oxide/cuteir/cute-rs/examples/global_store_compile.rs index 5b1ffd74f2..feb04b3fa9 100644 --- a/cuda-oxide/cuteir/cute-rs/examples/global_store_compile.rs +++ b/cuda-oxide/cuteir/cute-rs/examples/global_store_compile.rs @@ -59,4 +59,5 @@ pub unsafe fn store_global_f16(out: *mut f16, bits: u64) { } } +/// 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/cuteir/cute-rs/src/block_scaled.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled.rs deleted file mode 100644 index b3da2a61cb..0000000000 --- a/cuda-oxide/cuteir/cute-rs/src/block_scaled.rs +++ /dev/null @@ -1,860 +0,0 @@ -/* - * 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. - -use core::marker::PhantomData; - -use crate::cooperative::GmemMatrix; -use crate::markers::LeadingDim; -use crate::numeric::{E2M1, Mxf4E2M1, PackedE2M1x2, UE8M0, UE8M0x2, UE8M0x4}; -use crate::tiled_copy::{GlobalCopyTensor, SharedTensor}; -use crate::{Tensor, TensorElement, assume_div, load_tile}; - -/// 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; - -/// Rows whose K values are adjacent in memory. -/// -/// ```text -/// row 0: [k0 k1 k2 ...] -/// row 1: [k0 k1 k2 ...] -/// ``` -/// -/// `Mode` is either [`Mkl`] or [`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 { - rows: usize, - k: usize, - 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 { - rows: usize, - groups: usize, - 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, - } - } - - /// Return the byte offset for `(row, K group)`. - #[must_use] - #[inline(always)] - pub const fn offset(&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 - } -} - -/// The two [`UE8M0`] 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`]) separate from B scales ([`Nkl`]) so the MMA -/// call cannot swap them accidentally. -#[repr(transparent)] -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct MmaScalePair { - bits: u32, - 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 { - bits: u32, - 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, - } - } -} - -/// 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 Mxfp4BlockScaledTensor<'a, Mode: Copy> { - values: Tensor<'a, Mxf4E2M1, KMajor, f16>, - scales: Tensor<'a, UE8M0, DenseBlockScaleKMajor, u8>, -} - -impl<'a, Mode: Copy> Mxfp4BlockScaledTensor<'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.offset(row, pair * 2); - MmaScalePair { - bits: (self.scales.storage[offset] as u32) - | ((self.scales.storage[offset + 1] as u32) << 8), - mode: PhantomData, - } - } -} - -/// 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 { - rest_m: usize, - rest_k: usize, -} - -// Map one logical scale coordinate into Blackwell's 512-byte 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) - }}; -} - -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), - } - } - - /// Return the byte offset for `(batch, row, K group)`. - #[must_use] - #[inline(always)] - pub const fn offset(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 - } -} - -/// 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, -> { - values: Tensor<'a, Data, DataLayout, u8>, - scales: Tensor<'a, Scale, Sm1xxBlockScaleKMajor, u8>, -} - -impl<'a, Mode: Copy> BlockScaledTensor<'a, E2M1, UE8M0, 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) - } -} - -/// 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> { - values: &'a [u8], - scales: &'a [u8], - value_row_base: usize, - scale_row_base: usize, - 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> { - values: &'a [u8], - scales: &'a [u8], - value_base: usize, - scale_base: usize, - 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 { - values_lo: [u8; 16], - values_hi: [u8; 16], - scale0: f32, - scale1: f32, - scale2: f32, - scale3: f32, - mode: PhantomData, -} - -impl LoadedBlockScaledTile64 { - /// Return packed FP4 pair `PAIR` in increasing K order. - #[must_use] - #[inline(always)] - pub fn value_pair(self) -> PackedE2M1x2 { - const { assert!(PAIR < 32) }; - let bits = if PAIR < 16 { - self.values_lo[PAIR] - } else { - self.values_hi[PAIR - 16] - }; - PackedE2M1x2::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) - } -} - -/// Stable calls recognized by the CuTe importer. -/// -/// Kernel authors use the methods above. Keeping these functions in one -/// private module gives the compiler an exact boundary that does not depend -/// on whether rustc inlines an ergonomic method wrapper. -mod __compiler { - use super::*; - - #[inline(never)] - pub(super) const fn block_scaled_make<'a, Mode: Copy>( - values: &'a [u8], - scales: &'a [u8], - rows: usize, - k: usize, - ) -> BlockScaledTensor<'a, E2M1, UE8M0, 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)] - pub(super) fn block_scaled_thread_row<'a, Mode: Copy>( - tensor: BlockScaledTensor<'a, E2M1, UE8M0, 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)] - 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)] - 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_div::<16>(tile.value_base) }; - let scale_base = unsafe { assume_div::<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)] - pub(super) fn block_scaled_dot_k64( - lhs: LoadedBlockScaledTile64, - rhs: LoadedBlockScaledTile64, - mut acc: f32, - ) -> f32 { - 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 - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn canonical_scale_layout_matches_blackwell_golden_offsets() { - let layout = Sm1xxBlockScaleKMajor::<16>::new(256, 128); - - assert_eq!(layout.offset(0, 0, 0), 0); - assert_eq!(layout.offset(0, 31, 0), 496); - assert_eq!(layout.offset(0, 32, 0), 4); - assert_eq!(layout.offset(0, 127, 3), 511); - assert_eq!(layout.offset(0, 0, 4), 512); - assert_eq!(layout.offset(0, 128, 0), 1_024); - assert_eq!(layout.offset(1, 0, 0), 2_048); - assert_eq!(layout.storage_len(2), 4_096); - } - - #[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); - } - - #[test] - fn mxf4_tensor_keeps_logical_k_and_materializes_packed_carriers() { - let values = [f16::from_bits(0); 128]; - let scales = [127u8; 16]; - let matrix = Mxfp4BlockScaledTensor::::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); - } - - #[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); - - #[repr(align(4))] - struct AlignedScales([u8; 32]); - let values = [f16::from_bits(0); 256]; - let scales = AlignedScales([127; 32]); - let matrix = Mxfp4BlockScaledTensor::::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); - } - - #[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.offset(32, 0), 128); - assert_eq!(dense.storage_len(), 512); - assert_eq!(canonical.offset(0, 32, 0), 4); - assert_eq!(canonical.storage_len(1), 512); - } - - #[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.offset(0, row, 0)); - for group in 0..CANONICAL_SCALE_K_GROUPS { - assert_eq!(word * 4 + group, canonical.offset(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/cuteir/cute-rs/src/block_scaled/__compiler.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled/__compiler.rs new file mode 100644 index 0000000000..658738d7ae --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/block_scaled/__compiler.rs @@ -0,0 +1,132 @@ +/* + * 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::numeric::{E2M1, UE8M0, UE8M0x2}; +use crate::{Tensor, assume_div, load_tile}; +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, E2M1, UE8M0, 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, E2M1, UE8M0, 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_div::<16>(tile.value_base) }; + let scale_base = unsafe { assume_div::<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 +} diff --git a/cuda-oxide/cuteir/cute-rs/src/block_scaled/gemv.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled/gemv.rs new file mode 100644 index 0000000000..c8c6e5dec6 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/block_scaled/gemv.rs @@ -0,0 +1,150 @@ +/* + * 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::PackedE2M1x2; +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) -> PackedE2M1x2 { + const { assert!(PAIR < 32) }; + let bits = if PAIR < 16 { + self.values_lo[PAIR] + } else { + self.values_hi[PAIR - 16] + }; + PackedE2M1x2::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) + } +} diff --git a/cuda-oxide/cuteir/cute-rs/src/block_scaled/layouts.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled/layouts.rs new file mode 100644 index 0000000000..2564a33e28 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/block_scaled/layouts.rs @@ -0,0 +1,219 @@ +/* + * 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::numeric::UE8M0x4; +use crate::tiled_copy::SharedTensor; +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, + } + } + + /// Return the byte offset for `(row, K group)`. + #[must_use] + #[inline(always)] + pub const fn offset(&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), + } + } + + /// Return the byte offset for `(batch, row, K group)`. + #[must_use] + #[inline(always)] + pub const fn offset(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/cuteir/cute-rs/src/block_scaled/mod.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled/mod.rs new file mode 100644 index 0000000000..d084785f93 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/block_scaled/mod.rs @@ -0,0 +1,80 @@ +/* + * 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}; +pub use layouts::{ + CANONICAL_SCALE_ATOM_BYTES, CANONICAL_SCALE_K_GROUPS, CANONICAL_SCALE_M, DenseBlockScaleKMajor, + KMajor, SharedScaleAtom, Sm1xxBlockScaleKMajor, Sm120ScaleAtom, +}; +pub use scales::{MmaScalePair, ScalePack4}; +pub use tensor::{BlockScaledTensor, Mxfp4BlockScaledTensor}; + +/// 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/cuteir/cute-rs/src/block_scaled/scales.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled/scales.rs new file mode 100644 index 0000000000..13928efcaf --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/block_scaled/scales.rs @@ -0,0 +1,100 @@ +/* + * 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 [`UE8M0`](crate::numeric::UE8M0) 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, + } + } +} diff --git a/cuda-oxide/cuteir/cute-rs/src/block_scaled/tensor.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled/tensor.rs new file mode 100644 index 0000000000..901e6a2e4a --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/block_scaled/tensor.rs @@ -0,0 +1,170 @@ +/* + * 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, Sm1xxBlockScaleKMajor, +}; +use crate::cooperative::GmemMatrix; +use crate::markers::LeadingDim; +use crate::numeric::{E2M1, Mxf4E2M1, UE8M0, UE8M0x4}; +use crate::tiled_copy::GlobalCopyTensor; +use crate::{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 Mxfp4BlockScaledTensor<'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, UE8M0, DenseBlockScaleKMajor, u8>, +} + +impl<'a, Mode: Copy> Mxfp4BlockScaledTensor<'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.offset(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, E2M1, UE8M0, 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) + } +} diff --git a/cuda-oxide/cuteir/cute-rs/src/block_scaled/tests.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled/tests.rs new file mode 100644 index 0000000000..c2cf90f917 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/block_scaled/tests.rs @@ -0,0 +1,122 @@ +/* + * 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::{E2M1, UE8M0}; + +/// 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.offset(0, 0, 0), 0); + assert_eq!(layout.offset(0, 31, 0), 496); + assert_eq!(layout.offset(0, 32, 0), 4); + assert_eq!(layout.offset(0, 127, 3), 511); + assert_eq!(layout.offset(0, 0, 4), 512); + assert_eq!(layout.offset(0, 128, 0), 1_024); + assert_eq!(layout.offset(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 = Mxfp4BlockScaledTensor::::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 = Mxfp4BlockScaledTensor::::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.offset(32, 0), 128); + assert_eq!(dense.storage_len(), 512); + assert_eq!(canonical.offset(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.offset(0, row, 0)); + for group in 0..CANONICAL_SCALE_K_GROUPS { + assert_eq!(word * 4 + group, canonical.offset(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/cuteir/cute-rs/src/block_scaled_mma.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma.rs deleted file mode 100644 index e124affb0f..0000000000 --- a/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma.rs +++ /dev/null @@ -1,965 +0,0 @@ -/* - * 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. - -use core::marker::PhantomData; - -use crate::block_scaled::{Mkl, MmaScalePair, Nkl, ScalePack4, SharedScaleAtom, Sm120ScaleAtom}; -use crate::epilogue::Sm120EpilogueWarp128x128; -use crate::markers::ColMajor; -use crate::mma::{AccC, FragA, FragB, load_matrix_a}; -use crate::numeric::{Mxf4E2M1, UE8M0x4}; -use crate::tiled_copy::SharedTensor; - -/// 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 { - a: ScalePack4, - b0: ScalePack4, - b1: ScalePack4, - b2: ScalePack4, - 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 { - a0: ScalePack4, - a1: ScalePack4, - b0: ScalePack4, - b1: ScalePack4, - b2: ScalePack4, - b3: ScalePack4, - b4: ScalePack4, - b5: ScalePack4, - b6: ScalePack4, - 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 { - c00: AccC, - c01: AccC, - c02: AccC, - c03: AccC, - c04: AccC, - c05: AccC, - c06: AccC, - c07: AccC, - c10: AccC, - c11: AccC, - c12: AccC, - c13: AccC, - c14: AccC, - c15: AccC, - c16: AccC, - c17: AccC, -} - -/// 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::accumulate_k64`], 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> { - base: *mut f16, - capacity: usize, - warp_n: usize, - k_half: usize, - 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)] - unsafe fn load(&self, mma: &Mxfp4TiledMma) -> (FragB, FragB) { - 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. - /// - /// # 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 accumulate_k64( - &mut self, - mma: &Mxfp4TiledMma, - a: (FragA, FragA), - b: Mxf4BTileK64<'_, SharedLayout>, - scales: Mxf4ScalePairs128, - ) { - unsafe { __compiler::tiled_gemm(self, mma, a, b, scales) }; - } -} - -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) }; - } -} - -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) } - } -} - -// Map `(warp M position, first-or-second M block)` to a 16-row block number. -#[inline(always)] -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)] -const fn b_pair_atom_128(warp_n: usize, pair_n: usize) -> usize { - warp_n + 2 * pair_n -} - -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) }, - } - } -} - -/// 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 Mxfp4TiledMma { - lane: u32, - layout: PhantomData, -} - -impl Mxfp4TiledMma { - /// 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::accumulate_k64`]. `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), - ) -> FragA { - 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), - ) -> (FragB, FragB) { - let loaded = - unsafe { load_matrix_a::(source.carrier(), warp_tile, self.lane) }; - ( - FragB([loaded.0[0], loaded.0[2]]), - FragB([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, - ) -> (FragA, FragA) { - 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, - ) -> (FragB, FragB) { - 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`] - /// 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`] 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: FragA, - scale_a: MmaScalePair, - b: FragB, - scale_b: MmaScalePair, - c: AccC, - ) -> AccC { - let _ = (self, a, scale_a, b, scale_b, c); - unreachable!("cute-rs block-scaled MMA executed outside recognized CuTe device compilation") - } -} - -/// Stable compiler boundaries for the SM120 shared-memory MMA path. -/// -/// Public methods above keep the kernel readable. These exact calls let the -/// importer retain the same tensor, fragment, and tiled-MMA meaning whether -/// rustc inlines those wrappers or leaves them as calls. -#[doc(hidden)] -pub mod __compiler { - use super::*; - - #[inline(never)] - pub const fn tiled_mma_slice(lane: u32) -> Mxfp4TiledMma { - Mxfp4TiledMma { - lane, - layout: PhantomData, - } - } - - #[inline(never)] - pub const fn fragment_fill(fill: f32) -> Mxf4AccumulatorTile2x8 { - let cell = AccC([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, - } - } - - #[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); - } - } - - #[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) }, - } - } - - #[inline(never)] - pub unsafe fn mma_load_scales( - mma: &Mxfp4TiledMma, - 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))), - } - } - } - - #[inline(never)] - pub unsafe fn mma_load_a( - mma: &Mxfp4TiledMma, - source: &SharedTensor, - warp_m: usize, - k_half: usize, - ) -> (FragA, FragA) { - 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)), - ) - } - } - - #[inline(never)] - pub const fn mma_partition_b<'a, SharedLayout>( - _mma: &Mxfp4TiledMma, - 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, - } - } - - #[inline(never)] - pub unsafe fn tiled_gemm( - accumulators: &mut Mxf4AccumulatorTile2x8, - mma: &Mxfp4TiledMma, - a: (FragA, FragA), - 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) }; - } - } -} - -#[cfg(test)] -mod tests { - use super::{Mxf4AccumulatorTile2x8, Mxf4BTileK64, Mxfp4TiledMma, a_atom_128, b_pair_atom_128}; - use crate::block_scaled::{Mkl, Nkl, SharedScaleAtom, Sm120ScaleAtom}; - use crate::mma::AccC; - - #[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::() - ); - } - - #[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]] - ); - } - - #[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 = Mxfp4TiledMma::<()>::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 = Mxfp4TiledMma::<()>::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/cuteir/cute-rs/src/block_scaled_mma/__compiler.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/__compiler.rs new file mode 100644 index 0000000000..5977295b97 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/__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::epilogue::Sm120EpilogueWarp128x128; +use crate::mma::{AccC, FragA}; +use crate::numeric::Mxf4E2M1; +use crate::tiled_copy::SharedTensor; +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) -> Mxfp4TiledMma { + Mxfp4TiledMma { + 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 = AccC([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: &Mxfp4TiledMma, + 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 Mxfp4TiledMma::load_a. +#[inline(never)] +pub unsafe fn mma_load_a( + mma: &Mxfp4TiledMma, + source: &SharedTensor, + warp_m: usize, + k_half: usize, +) -> (FragA, FragA) { + 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: &Mxfp4TiledMma, + 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: &Mxfp4TiledMma, + a: (FragA, FragA), + 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/cuteir/cute-rs/src/block_scaled_mma/epilogue.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/epilogue.rs new file mode 100644 index 0000000000..84f299ab05 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/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::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/cuteir/cute-rs/src/block_scaled_mma/fragments.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/fragments.rs new file mode 100644 index 0000000000..2876cda942 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/fragments.rs @@ -0,0 +1,308 @@ +/* + * 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, Mxfp4TiledMma}; +use crate::block_scaled::{Mkl, MmaScalePair, Nkl, ScalePack4}; +use crate::mma::{AccC, FragA, FragB}; +use crate::numeric::Mxf4E2M1; +use crate::tiled_copy::SharedTensor; +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: AccC, + /// Per-lane accumulators for row band 0 and N8 fragment 1. + pub(super) c01: AccC, + /// Per-lane accumulators for row band 0 and N8 fragment 2. + pub(super) c02: AccC, + /// Per-lane accumulators for row band 0 and N8 fragment 3. + pub(super) c03: AccC, + /// Per-lane accumulators for row band 0 and N8 fragment 4. + pub(super) c04: AccC, + /// Per-lane accumulators for row band 0 and N8 fragment 5. + pub(super) c05: AccC, + /// Per-lane accumulators for row band 0 and N8 fragment 6. + pub(super) c06: AccC, + /// Per-lane accumulators for row band 0 and N8 fragment 7. + pub(super) c07: AccC, + /// Per-lane accumulators for row band 1 and N8 fragment 0. + pub(super) c10: AccC, + /// Per-lane accumulators for row band 1 and N8 fragment 1. + pub(super) c11: AccC, + /// Per-lane accumulators for row band 1 and N8 fragment 2. + pub(super) c12: AccC, + /// Per-lane accumulators for row band 1 and N8 fragment 3. + pub(super) c13: AccC, + /// Per-lane accumulators for row band 1 and N8 fragment 4. + pub(super) c14: AccC, + /// Per-lane accumulators for row band 1 and N8 fragment 5. + pub(super) c15: AccC, + /// Per-lane accumulators for row band 1 and N8 fragment 6. + pub(super) c16: AccC, + /// Per-lane accumulators for row band 1 and N8 fragment 7. + pub(super) c17: AccC, +} + +/// 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::accumulate_k64`], 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: &Mxfp4TiledMma, + ) -> (FragB, FragB) { + 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. + /// + /// # 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 accumulate_k64( + &mut self, + mma: &Mxfp4TiledMma, + a: (FragA, FragA), + 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/cuteir/cute-rs/src/block_scaled_mma/mod.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/mod.rs new file mode 100644 index 0000000000..cdc75eabd5 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/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::Mxfp4TiledMma; + +/// 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/cuteir/cute-rs/src/block_scaled_mma/tests.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/tests.rs new file mode 100644 index 0000000000..51f8c8f4ae --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/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, Mxfp4TiledMma}; +use crate::block_scaled::{Mkl, Nkl, SharedScaleAtom, Sm120ScaleAtom}; +use crate::mma::AccC; + +/// 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 = Mxfp4TiledMma::<()>::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 = Mxfp4TiledMma::<()>::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/cuteir/cute-rs/src/block_scaled_mma/tiled.rs b/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/tiled.rs new file mode 100644 index 0000000000..a41b081c3a --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/block_scaled_mma/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::markers::ColMajor; +use crate::mma::{AccC, FragA, FragB, load_matrix_a}; +use crate::numeric::{Mxf4E2M1, UE8M0x4}; +use crate::tiled_copy::SharedTensor; +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 Mxfp4TiledMma { + /// 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 Mxfp4TiledMma { + /// 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::accumulate_k64`](super::Mxf4AccumulatorTile2x8::accumulate_k64). `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), + ) -> FragA { + 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), + ) -> (FragB, FragB) { + let loaded = + unsafe { load_matrix_a::(source.carrier(), warp_tile, self.lane) }; + ( + FragB([loaded.0[0], loaded.0[2]]), + FragB([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, + ) -> (FragA, FragA) { + 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, + ) -> (FragB, FragB) { + 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: FragA, + scale_a: MmaScalePair, + b: FragB, + scale_b: MmaScalePair, + c: AccC, + ) -> AccC { + 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/cuteir/cute-rs/src/cooperative.rs b/cuda-oxide/cuteir/cute-rs/src/cooperative.rs index 43e9cf591d..545f16e605 100644 --- a/cuda-oxide/cuteir/cute-rs/src/cooperative.rs +++ b/cuda-oxide/cuteir/cute-rs/src/cooperative.rs @@ -120,10 +120,12 @@ pub unsafe fn copy_g2s( 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: diff --git a/cuda-oxide/cuteir/cute-rs/src/epilogue.rs b/cuda-oxide/cuteir/cute-rs/src/epilogue.rs index 4df6fa00c7..148caa6da9 100644 --- a/cuda-oxide/cuteir/cute-rs/src/epilogue.rs +++ b/cuda-oxide/cuteir/cute-rs/src/epilogue.rs @@ -72,6 +72,7 @@ pub const SM120_EPILOGUE_ELEMENTS: usize = SM120_EPILOGUE_ROWS * SM120_EPILOGUE_ /// 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. @@ -124,6 +125,7 @@ pub const fn sm120_epilogue_physical_offset(row: usize, col: usize) -> usize { #[repr(transparent)] #[derive(Clone, Copy)] pub struct Sm120Epilogue128x128 { + /// Base of the complete 32-KiB FP16 shared output tile. base: *mut f16, } @@ -212,11 +214,13 @@ impl Sm120Epilogue128x128 { 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, @@ -230,18 +234,27 @@ pub mod __compiler { } } + /// 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, @@ -258,8 +271,11 @@ pub mod __compiler { /// 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, } @@ -295,12 +311,14 @@ impl Sm120EpilogueWarp128x128 { } } +/// 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!(::ROWS, 128); @@ -314,6 +332,7 @@ mod tests { 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]; @@ -341,6 +360,7 @@ mod tests { 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]; @@ -357,6 +377,7 @@ mod tests { 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 { diff --git a/cuda-oxide/cuteir/cute-rs/src/lib.rs b/cuda-oxide/cuteir/cute-rs/src/lib.rs index 60ddf54b8f..2b80ecb748 100644 --- a/cuda-oxide/cuteir/cute-rs/src/lib.rs +++ b/cuda-oxide/cuteir/cute-rs/src/lib.rs @@ -5,21 +5,21 @@ //! CuTe-style building blocks for GPU kernels written in Rust. //! -//! The types describe what a kernel means. The compiler turns a small set of -//! recognized functions into GPU operations: +//! A tensor view describes storage without loading it. Split it into tiles, +//! then select the part one thread will process: //! -//! ```text -//! Tensor + Layout + TiledCopy -//! │ -//! ▼ -//! Rust kernel code -//! │ -//! ▼ -//! load / copy / MMA / store -//! │ -//! ▼ -//! selected GPU backend //! ``` +//! 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/cuteir/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. //! //! 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 diff --git a/cuda-oxide/cuteir/cute-rs/src/markers.rs b/cuda-oxide/cuteir/cute-rs/src/markers.rs index 5db9f90c2f..3f4d1a343b 100644 --- a/cuda-oxide/cuteir/cute-rs/src/markers.rs +++ b/cuda-oxide/cuteir/cute-rs/src/markers.rs @@ -29,13 +29,17 @@ pub struct C; /// /// Pairs can be nested. For example, `(A, B), C` and `A, (B, C)` describe /// different layouts. -pub struct T2(PhantomData<(A, B)>); +pub struct T2( + #[doc = "Carries the two nested compile-time tuple entries."] PhantomData<(A, B)>, +); /// A layout made from a `Shape` and its `Stride`. /// /// A stride is the number of elements to skip when a coordinate increases by /// one. -pub struct L(PhantomData<(Shape, Stride)>); +pub struct L( + #[doc = "Carries corresponding shape and stride descriptions."] PhantomData<(Shape, Stride)>, +); /// A `ROWS × COLS` matrix stored one complete row after another. /// @@ -113,7 +117,9 @@ pub struct Swizzle; /// /// `OFFSET` and `Inner` count elements here. The compiler converts the whole /// rule to bytes after it knows the element size. -pub struct Composed(PhantomData<(Outer, Inner)>); +pub struct Composed( + #[doc = "Carries the inner map and outer offset permutation."] PhantomData<(Outer, Inner)>, +); /// Makes a two-dimensional shared-memory layout readable by host code. /// @@ -195,15 +201,21 @@ impl, C<8>>; + /// Row-major marker strides corresponding to the test shape. type Stride = T2, C<1>>; + /// Pairs the test’s shape and strides without runtime storage. type Affine = L; + /// Adds a swizzle to verify composed markers also remain zero-sized. type Smem = Composed, 0, Affine>; assert_eq!(core::mem::size_of::>(), 0); diff --git a/cuda-oxide/cuteir/cute-rs/src/mma.rs b/cuda-oxide/cuteir/cute-rs/src/mma.rs index f9b7feab97..98c7853d0c 100644 --- a/cuda-oxide/cuteir/cute-rs/src/mma.rs +++ b/cuda-oxide/cuteir/cute-rs/src/mma.rs @@ -33,12 +33,12 @@ use crate::SmemTile; /// order expected by `mma.sync`. #[repr(transparent)] #[derive(Clone, Copy)] -pub struct FragA(pub [u32; 4]); +pub struct FragA(#[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 FragB(pub [u32; 2]); +pub struct FragB(#[doc = "This lane’s two packed B registers in instruction order."] pub [u32; 2]); /// Two adjacent 16x8 B fragments in MMA register order. /// @@ -46,7 +46,7 @@ pub struct FragB(pub [u32; 2]); /// retiling an A-fragment load with [`FragA::as_b_pair`]. #[repr(transparent)] #[derive(Clone, Copy)] -pub struct FragBPair(pub [u32; 4]); +pub struct FragBPair(#[doc = "Packed B registers for two neighboring N8 fragments."] pub [u32; 4]); impl FragBPair { /// Select the left (`N = 0`) or right (`N = 1`) fragment. @@ -72,7 +72,9 @@ impl FragA { /// `cuda_device::mma_frag` maps these four values to result coordinates. #[repr(transparent)] #[derive(Clone, Copy)] -pub struct AccC(pub [f32; 4]); +pub struct AccC( + #[doc = "This lane’s four FP32 result registers in instruction order."] pub [f32; 4], +); impl AccC { /// Four zero values for starting `C = A*B + C`. @@ -85,7 +87,18 @@ impl AccC { #[must_use = "use the returned accumulator; mma does not modify self in place"] #[inline(always)] pub unsafe fn mma(self, a: FragA, b: FragB) -> Self { - Self(unsafe { cuda_device::wmma::mma_m16n8k16_f32_f16(self.0, a.0, b.0) }) + 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(), + ) } } diff --git a/cuda-oxide/cuteir/cute-rs/src/numeric.rs b/cuda-oxide/cuteir/cute-rs/src/numeric.rs index 3092cadd8d..dd5014b9ba 100644 --- a/cuda-oxide/cuteir/cute-rs/src/numeric.rs +++ b/cuda-oxide/cuteir/cute-rs/src/numeric.rs @@ -91,7 +91,7 @@ impl TensorElement for UE8M0x4 {} /// Every bit pattern is valid, including positive and negative zero. #[repr(transparent)] #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct PackedE2M1x2(u8); +pub struct PackedE2M1x2(#[doc = "Low and high FP4 nibbles preserved in one byte."] u8); impl PackedE2M1x2 { /// Treat `bits` as two packed E2M1 values without changing them. @@ -132,7 +132,7 @@ impl PackedE2M1x2 { /// Every byte is valid. The byte `255` represents NaN. #[repr(transparent)] #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct UE8M0x2(u16); +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. @@ -174,7 +174,9 @@ impl UE8M0x2 { } #[inline(always)] +/// Decode the low FP4 E2M1 nibble, including its sign bit, to f32. fn e2m1_to_f32(nibble: u8) -> f32 { + /// Exact positive E2M1 values indexed by the three magnitude bits. const MAGNITUDES: [f32; 8] = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]; let magnitude = MAGNITUDES[(nibble & 0x07) as usize]; if nibble & 0x08 == 0 { @@ -185,6 +187,7 @@ fn e2m1_to_f32(nibble: u8) -> f32 { } #[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 @@ -197,10 +200,12 @@ fn ue8m0_to_f32(byte: u8) -> f32 { } } +/// 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 { @@ -208,6 +213,7 @@ mod tests { } } + /// 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); @@ -215,8 +221,10 @@ mod tests { 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()); @@ -227,6 +235,7 @@ mod tests { } } + /// 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); diff --git a/cuda-oxide/cuteir/cute-rs/src/pipeline.rs b/cuda-oxide/cuteir/cute-rs/src/pipeline.rs index fb41a6cb49..2552cc86b0 100644 --- a/cuda-oxide/cuteir/cute-rs/src/pipeline.rs +++ b/cuda-oxide/cuteir/cute-rs/src/pipeline.rs @@ -49,8 +49,11 @@ pub enum Consumer {} /// 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, } @@ -119,7 +122,9 @@ impl PipelineState { /// buffers in shared memory. #[derive(Clone, Copy)] pub struct TmaLoadPipeline { + /// 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, } @@ -350,24 +355,28 @@ impl 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 _ = TmaStorePipeline::::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 _ = TmaStorePipeline::::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 _ = TmaStorePipeline::::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) }; @@ -379,6 +388,7 @@ pub mod __compiler { } } + /// Construct slot zero with full parity zero, validating a nonempty ring. #[inline(never)] pub fn pipeline_state_new_consumer() -> PipelineState { const { assert!(STAGES > 0) }; @@ -390,6 +400,7 @@ pub mod __compiler { } } + /// Compiler boundary exposing the current reusable stage index. #[inline(never)] pub fn pipeline_state_slot( state: &PipelineState, @@ -397,6 +408,7 @@ pub mod __compiler { 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, @@ -410,6 +422,10 @@ pub mod __compiler { } } + /// Compiler boundary splitting adjacent full and empty barrier arrays. + /// + /// # Safety + /// The base must satisfy the public TmaLoadPipeline storage and lifetime contract. #[inline(never)] pub unsafe fn tma_load_pipeline_from_raw_base< const STAGES: usize, @@ -425,6 +441,10 @@ pub mod __compiler { } } + /// 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, @@ -439,6 +459,10 @@ pub mod __compiler { 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, @@ -453,6 +477,10 @@ pub mod __compiler { 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, @@ -467,6 +495,10 @@ pub mod __compiler { 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, @@ -481,6 +513,10 @@ pub mod __compiler { 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, @@ -495,6 +531,10 @@ pub mod __compiler { 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, @@ -510,10 +550,12 @@ pub mod __compiler { } } +/// 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(); @@ -522,6 +564,7 @@ mod tests { assert_eq!((consumer.slot(), 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(); @@ -532,20 +575,25 @@ mod tests { } } + /// 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 = TmaLoadPipeline<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); diff --git a/cuda-oxide/cuteir/cute-rs/src/scheduler.rs b/cuda-oxide/cuteir/cute-rs/src/scheduler.rs index 9c68b468e5..54a9257902 100644 --- a/cuda-oxide/cuteir/cute-rs/src/scheduler.rs +++ b/cuda-oxide/cuteir/cute-rs/src/scheduler.rs @@ -41,9 +41,13 @@ /// the selected tile. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct WorkTile { + /// 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, } @@ -110,7 +114,9 @@ pub struct StaticPersistentTileScheduler< 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, } @@ -251,6 +257,7 @@ impl pub mod __compiler { use super::{StaticPersistentTileScheduler, WorkTile}; + /// Device-only boundary reading the block index and grid stride for a persistent schedule. #[inline(never)] pub fn scheduler_new_1d() -> StaticPersistentTileScheduler { @@ -258,6 +265,7 @@ pub mod __compiler { 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, @@ -266,6 +274,7 @@ pub mod __compiler { < 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, @@ -287,6 +296,7 @@ pub mod __compiler { } } + /// Expose the flattened and three logical coordinates without recomputing them. #[inline(never)] pub fn work_tile_coordinates< const M_TILES: usize, @@ -298,6 +308,7 @@ pub mod __compiler { (tile.linear, 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, @@ -306,13 +317,16 @@ pub mod __compiler { } } +/// 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), @@ -334,9 +348,12 @@ mod tests { } } + /// 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]; @@ -357,8 +374,10 @@ mod tests { 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 { @@ -372,16 +391,20 @@ mod tests { } } + /// 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()); } + /// Resident-grid sizing must cap at total work and reject arithmetic or launch-width overflow. #[test] fn resident_grid_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::resident_grid_size(170, 1), Some(170)); assert_eq!(Scheduler::resident_grid_size(170, 2), Some(256)); @@ -389,12 +412,15 @@ mod tests { assert_eq!(Scheduler::resident_grid_size(0, 1), None); assert_eq!(Scheduler::resident_grid_size(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::resident_grid_size(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.current().is_some()); @@ -404,8 +430,10 @@ mod tests { assert!(scheduler.current().is_none()); } + /// 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/cuteir/cute-rs/src/sm100.rs b/cuda-oxide/cuteir/cute-rs/src/sm100.rs deleted file mode 100644 index 662deaca7b..0000000000 --- a/cuda-oxide/cuteir/cute-rs/src/sm100.rs +++ /dev/null @@ -1,522 +0,0 @@ -/* - * 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 -//! ``` - -use crate::TmaDesc; -use crate::markers::ReifySmem2D; -use core::marker::PhantomData; -use cuda_device::barrier::Barrier; - -/// A shared tile used as the left input of a matrix product. -pub enum OperandA {} -/// A shared tile used as the right input of a matrix product. -pub enum OperandB {} - -/// A non-owning FP16 shared tile with a layout and an operand role. -pub struct Sm100SharedTile { - base: *mut f16, - 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, - } - } -} - -/// 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 { - 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 allocate(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 deallocate(address: u32) { - let () = Self::VALID; - unsafe { __compiler::sm100_tmem_dealloc::(address) } - } -} - -/// 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. -pub struct Sm100TiledMma( - PhantomData<(AL, B0L, B1L)>, -); -impl< - AL: ReifySmem2D, - B0L: ReifySmem2D, - B1L: ReifySmem2D, - const N0: u32, - const N1: u32, - const K: u32, -> Sm100TiledMma -{ - 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::ROWS == 128 && AL::COLS == K as i64); - assert!(B0L::ROWS == (N0 / 2) as i64 && B0L::COLS == K as i64); - assert!(B1L::ROWS == (N1 / 2) as i64 && B1L::COLS == 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() - } -} - -/// A cluster TMA copy whose descriptor and destination have the same layout. -pub struct Sm100ClusterTmaCopy(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 TmaDesc, - 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) - } - } -} - -/// Convert one 128x32 FP32 accumulator slice to an FP16 tile in shared memory. -pub struct Sm100TmemEpilogue(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 TmaDesc, - src: *const f16, - row: u32, - column: u32, - ) { - unsafe { __compiler::sm100_tma_store::(desc, src, row, column) } - } -} - -/// 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 Sm100TmaMmaPipeline { - full: *mut Barrier, - empty: *mut Barrier, -} -impl Sm100TmaMmaPipeline { - 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(&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 full_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 Sm100AccumulatorPipeline { - empty: *mut Barrier, - full: *mut Barrier, -} -impl Sm100AccumulatorPipeline { - 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. -#[derive(Clone, Copy, Default)] -pub struct Sm100TmaStorePipeline; -impl Sm100TmaStorePipeline { - 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::() } - } -} - -/// Calls recognized by the compiler; never execute these on the host. -#[doc(hidden)] -#[allow(missing_docs)] -pub mod __compiler { - use super::*; - #[inline(never)] - pub unsafe fn sm100_tmem_alloc(token: *mut u32) { - let _ = token; - unreachable!("SM100 CuTe marker requires device compilation") - } - #[inline(never)] - pub unsafe fn sm100_tmem_dealloc(tmem: u32) { - let _ = tmem; - unreachable!("SM100 CuTe marker requires device compilation") - } - #[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") - } - #[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. - #[inline(never)] - pub unsafe fn sm100_cluster_tma_load( - desc: *const TmaDesc, - 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") - } - - #[inline(never)] - pub unsafe fn sm100_tma_store( - desc: *const TmaDesc, - src: *const f16, - row: u32, - col: u32, - ) { - let _ = (desc, src, row, col); - unreachable!("SM100 CuTe marker requires device compilation") - } - macro_rules! store_marker { - ($name:ident) => { - #[inline(never)] - pub unsafe fn $name() { - unreachable!("SM100 CuTe marker requires device compilation") - } - }; - } - store_marker!(sm100_store_commit); - store_marker!(sm100_store_acquire); - store_marker!(sm100_store_tail); - macro_rules! pipeline_marker { - ($name:ident ($($arg:ident:$ty:ty),*)) => { - #[inline(never)] - pub unsafe fn $name($($arg:$ty),*) { - let _ = ($($arg),*); unreachable!("SM100 CuTe marker requires device compilation") - } - }; - } - pipeline_marker!(sm100_pipeline_init(full: *mut Barrier, empty: *mut Barrier)); - pipeline_marker!(sm100_pipeline_acquire(empty: *mut Barrier, stage: u32, phase: u32)); - pipeline_marker!(sm100_pipeline_expect(full: *mut Barrier, stage: u32)); - pipeline_marker!(sm100_pipeline_wait(full: *mut Barrier, stage: u32, phase: u32)); - pipeline_marker!(sm100_pipeline_release(empty: *mut Barrier, stage: u32)); - macro_rules! accumulator_marker { - ($name:ident ($($arg:ident:$ty:ty),*)) => { - #[inline(never)] - pub unsafe fn $name($($arg:$ty),*) { - let _ = ($($arg),*); unreachable!("SM100 CuTe marker requires device compilation") - } - }; - } - accumulator_marker!(sm100_accumulator_init(empty: *mut Barrier, full: *mut Barrier)); - accumulator_marker!(sm100_accumulator_acquire(empty: *mut Barrier, phase: u32)); - accumulator_marker!(sm100_accumulator_commit(full: *mut Barrier)); - accumulator_marker!(sm100_accumulator_wait(full: *mut Barrier, phase: u32)); - accumulator_marker!(sm100_accumulator_release(empty: *mut Barrier)); -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::{Composed, RowMajor, Swizzle}; - type A = Composed, 0, RowMajor<128, 64>>; - type B0 = Composed, 0, RowMajor<96, 64>>; - type B1 = Composed, 0, RowMajor<80, 64>>; - #[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/cuteir/cute-rs/src/sm100/__compiler.rs b/cuda-oxide/cuteir/cute-rs/src/sm100/__compiler.rs new file mode 100644 index 0000000000..974fe55455 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/sm100/__compiler.rs @@ -0,0 +1,137 @@ +/* + * 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::TmaDesc; +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. +#[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. +#[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 TmaDesc, + 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 TmaDesc, + 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/cuteir/cute-rs/src/sm100/mma.rs b/cuda-oxide/cuteir/cute-rs/src/sm100/mma.rs new file mode 100644 index 0000000000..037ffa57ad --- /dev/null +++ b/cuda-oxide/cuteir/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::{Composed, RowMajor, Sm100TiledMma, Swizzle}; +/// type Tile = Composed, 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::ROWS == 128 && AL::COLS == K as i64); + assert!(B0L::ROWS == (N0 / 2) as i64 && B0L::COLS == K as i64); + assert!(B1L::ROWS == (N1 / 2) as i64 && B1L::COLS == 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/cuteir/cute-rs/src/sm100/mod.rs b/cuda-oxide/cuteir/cute-rs/src/sm100/mod.rs new file mode 100644 index 0000000000..2f36550304 --- /dev/null +++ b/cuda-oxide/cuteir/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 +//! ``` + +/// A shared tile used as the left input of a matrix product. +pub enum OperandA {} +/// A shared tile used as the right input of a matrix product. +pub enum OperandB {} + +mod mma; +mod pipeline; +mod shared; +mod tma; +mod tmem; + +pub use mma::Sm100TiledMma; +pub use pipeline::{Sm100AccumulatorPipeline, Sm100TmaMmaPipeline, Sm100TmaStorePipeline}; +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/cuteir/cute-rs/src/sm100/pipeline.rs b/cuda-oxide/cuteir/cute-rs/src/sm100/pipeline.rs new file mode 100644 index 0000000000..596d16d0dd --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/sm100/pipeline.rs @@ -0,0 +1,190 @@ +/* + * 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 Sm100TmaMmaPipeline { + /// 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 Sm100TmaMmaPipeline { + /// 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(&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 full_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 Sm100AccumulatorPipeline { + /// 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 Sm100AccumulatorPipeline { + /// 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. +#[derive(Clone, Copy, Default)] +pub struct Sm100TmaStorePipeline; +impl Sm100TmaStorePipeline { + /// 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/cuteir/cute-rs/src/sm100/shared.rs b/cuda-oxide/cuteir/cute-rs/src/sm100/shared.rs new file mode 100644 index 0000000000..0caa688efe --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/cute-rs/src/sm100/tests.rs b/cuda-oxide/cuteir/cute-rs/src/sm100/tests.rs new file mode 100644 index 0000000000..244ada585e --- /dev/null +++ b/cuda-oxide/cuteir/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::{Composed, RowMajor, Swizzle}; +/// Local 128-by-64 FP16 A tile with the B128 swizzle. +type A = Composed, 0, RowMajor<128, 64>>; +/// Local 96-by-64 B tile belonging to the collective N192 partition. +type B0 = Composed, 0, RowMajor<96, 64>>; +/// Local 80-by-64 B tile belonging to the collective N160 partition. +type B1 = Composed, 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/cuteir/cute-rs/src/sm100/tma.rs b/cuda-oxide/cuteir/cute-rs/src/sm100/tma.rs new file mode 100644 index 0000000000..079cb5a19c --- /dev/null +++ b/cuda-oxide/cuteir/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::TmaDesc; +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 TmaDesc, + 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/cuteir/cute-rs/src/sm100/tmem.rs b/cuda-oxide/cuteir/cute-rs/src/sm100/tmem.rs new file mode 100644 index 0000000000..3cf0900c34 --- /dev/null +++ b/cuda-oxide/cuteir/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::TmaDesc; +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 allocate(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 deallocate(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 TmaDesc, + src: *const f16, + row: u32, + column: u32, + ) { + unsafe { __compiler::sm100_tma_store::(desc, src, row, column) } + } +} diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention.rs deleted file mode 100644 index 83ecacff87..0000000000 --- a/cuda-oxide/cuteir/cute-rs/src/sm100_attention.rs +++ /dev/null @@ -1,536 +0,0 @@ -/* - * 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. -pub mod pipeline; - -use crate::{OperandA, OperandB, TmaDesc}; -use core::marker::PhantomData; -use cuda_device::{Barrier, CuSimd}; - -/// Physical layout of a shared FP16 operand, with CuTe's tiled byte swizzle. -mod sealed { - 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. - fn offset(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. -pub struct SharedLayout; -impl SharedLayout { - 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. - #[inline(always)] - pub fn offset(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 offset(row: u32, col: u32) -> usize { - Self::offset(row, col) - } -} - -/// Non-owning shared tile, typed by its physical layout and MMA operand role. -pub struct SharedTile { - base: *mut f16, - byte_offset: u32, - 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 { - base: *mut f16, - marker: PhantomData<(L, Role)>, -} -impl SharedTileArray { - 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 - }; - 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, - } - } -} - -/// Independent shared/shared FP16 MMA with FP32 accumulation. -pub struct MmaSs( - PhantomData<(AL, BL)>, -); -impl - MmaSs -{ - 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. -/// The importer currently supports K=32 or a positive multiple of 64. -pub struct MmaTs(PhantomData); -impl MmaTs { - 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() - } -} - -/// 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], -} -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) }; - cuda_device::tcgen05::tcgen05_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 { cuda_device::tcgen05::$load(address) }.to_array(); - let mut values = [0.0; $n]; - let mut i = 0; - while i < $n { - cuda_device::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) }; - cuda_device::tcgen05::tcgen05_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 { - cuda_device::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) }; - cuda_device::tcgen05::tcgen05_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 { cuda_device::tcgen05::$store(address, CuSimd::new(raw)) }; - } - } - }; -} -fragment!(2, tcgen05_ld_32x32b_x2_raw, tcgen05_st_32x32b_x2_raw); -fragment!(4, tcgen05_ld_32x32b_x4_raw, tcgen05_st_32x32b_x4_raw); -fragment!(8, tcgen05_ld_32x32b_x8_raw, tcgen05_st_32x32b_x8_raw); -fragment!(16, tcgen05_ld_32x32b_x16_raw, tcgen05_st_32x32b_x16_raw); -fragment!(32, tcgen05_ld_32x32b_x32_raw, tcgen05_st_32x32b_x32_raw); -// 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. -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) }; - cuda_device::tcgen05::tcgen05_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 { cuda_device::tcgen05::tcgen05_ld_32x32b_x32_raw(address + $offset) };)+ - let mut values = [0.0; $n]; - $( - let raw = $chunk.to_array(); - let mut i = 0; - while i < 32 { - cuda_device::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) }; - cuda_device::tcgen05::tcgen05_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 { - cuda_device::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) }; - cuda_device::tcgen05::tcgen05_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 { - cuda_device::thread::__unroll_config::<0>(); - chunk[i] = raw[$offset + i]; - i += 1; - } - unsafe { cuda_device::tcgen05::tcgen05_st_32x32b_x32_raw(address + $offset, CuSimd::new(chunk)) }; - )+ - } - } - } -} -chunked_fragment!(64; 0 => a, 32 => b); -chunked_fragment!(128; 0 => a, 32 => b, 64 => c, 96 => d); - -/// A rank-three TMA copy into the same physical layout used by MMA. -pub struct TmaCopy(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: &TmaDesc, - 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: &TmaDesc, 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, - ) - } - } -} - -#[doc(hidden)] -pub mod __compiler { - use super::*; - #[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") - } - #[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") - } - #[inline(never)] - #[allow(clippy::too_many_arguments)] - pub unsafe fn tma_load( - desc: *const TmaDesc, - 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") - } - #[inline(never)] - pub unsafe fn tma_store( - desc: *const TmaDesc, - src: *mut f16, - c0: i32, - c1: i32, - c2: i32, - ) { - let _ = (desc, src, c0, c1, c2); - unreachable!("CuTe compiler marker") - } -} - -#[cfg(feature = "host")] -mod host; -#[cfg(feature = "host")] -pub use host::make_tma_desc; diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/__compiler.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/__compiler.rs new file mode 100644 index 0000000000..9dbc602604 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/__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::TmaDesc; +use cuda_device::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 TmaDesc, + 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 TmaDesc, + src: *mut f16, + c0: i32, + c1: i32, + c2: i32, +) { + let _ = (desc, src, c0, c1, c2); + unreachable!("CuTe compiler marker") +} diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/host.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/host.rs index be592f9209..9fbc287a60 100644 --- a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/host.rs +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/host.rs @@ -14,7 +14,8 @@ use cuda_core::sys as driver; /// Large contiguous dimensions are split into swizzle-sized chunks. For an /// MN-major operand the logical rows/columns are transposed relative to global /// physical rows/columns. Global coordinates always use the physical ordering. -/// The descriptor borrows the allocation; callers retain it until completion. +/// 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 make_tma_desc( base: *mut c_void, @@ -42,7 +43,7 @@ pub fn make_tma_desc( { 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) % 16 != 0 { + 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 diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/mma.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/mma.rs new file mode 100644 index 0000000000..89b5848d04 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/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::{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::sm100_attention::{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::sm100_attention::{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/cuteir/cute-rs/src/sm100_attention/mod.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/mod.rs new file mode 100644 index 0000000000..d573b5e3b2 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/mod.rs @@ -0,0 +1,28 @@ +/* + * 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. +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::make_tma_desc; diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/pipeline.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/pipeline.rs index 5a219f5299..ab9d5b5071 100644 --- a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/pipeline.rs +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/pipeline.rs @@ -8,21 +8,42 @@ //! 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::{cluster, groups::CtaGroup2, memory::SharedAddress, tcgen05}; use cuda_device::barrier::{self, Barrier}; +/// 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 {} } @@ -193,6 +214,7 @@ 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; } @@ -209,13 +231,25 @@ impl CursorRole for Consumer { /// 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::sm100_attention::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. @@ -268,14 +302,66 @@ impl Default for RingCursor { /// 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, TmaDesc, sm100_attention::{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: &TmaDesc, b: &mut SharedTile, rank: u32) { +/// unsafe { +/// ring.acquire_cursor(*p); +/// if rank == 0 && ctm_rs::warp::is_elected_sync(u32::MAX) { +/// ring.expect_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, sm100_attention::{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::(); @@ -303,14 +389,17 @@ impl TwoCtaRing 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 } @@ -520,6 +609,7 @@ impl TwoCtaRing() { if ::TC_AFTER { tcgen05::fence_after_thread_sync(); @@ -527,12 +617,20 @@ fn after_wait() { } #[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::map_shared(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( @@ -542,11 +640,14 @@ unsafe fn mma_arrive(bar: *mut Barrier) { } } +/// 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), @@ -562,8 +663,10 @@ mod tests { 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)] { @@ -572,8 +675,10 @@ mod tests { 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); diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/shared.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/shared.rs new file mode 100644 index 0000000000..4419fd4947 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/shared.rs @@ -0,0 +1,168 @@ +/* + * 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. + fn offset(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::sm100_attention::SharedLayout; +/// type Tile = SharedLayout<64, 64>; +/// assert_eq!(Tile::offset(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. + #[inline(always)] + pub fn offset(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 offset(row: u32, col: u32) -> usize { + Self::offset(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/cuteir/cute-rs/src/sm100_attention/tma.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/tma.rs new file mode 100644 index 0000000000..e67f6afb21 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/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::TmaDesc; +use core::marker::PhantomData; +use cuda_device::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: &TmaDesc, + 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: &TmaDesc, 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/cuteir/cute-rs/src/sm100_attention/tmem.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/tmem.rs new file mode 100644 index 0000000000..d5ae2be419 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/tmem.rs @@ -0,0 +1,189 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Warp-collective register fragments for tensor-memory loads and stores. + +use cuda_device::CuSimd; + +/// 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) }; + cuda_device::tcgen05::tcgen05_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 { cuda_device::tcgen05::$load(address) }.to_array(); + let mut values = [0.0; $n]; + let mut i = 0; + while i < $n { + cuda_device::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) }; + cuda_device::tcgen05::tcgen05_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 { + cuda_device::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) }; + cuda_device::tcgen05::tcgen05_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 { cuda_device::tcgen05::$store(address, CuSimd::new(raw)) }; + } + } + }; +} +fragment!(2, tcgen05_ld_32x32b_x2_raw, tcgen05_st_32x32b_x2_raw); +fragment!(4, tcgen05_ld_32x32b_x4_raw, tcgen05_st_32x32b_x4_raw); +fragment!(8, tcgen05_ld_32x32b_x8_raw, tcgen05_st_32x32b_x8_raw); +fragment!(16, tcgen05_ld_32x32b_x16_raw, tcgen05_st_32x32b_x16_raw); +fragment!(32, tcgen05_ld_32x32b_x32_raw, tcgen05_st_32x32b_x32_raw); +// 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) }; + cuda_device::tcgen05::tcgen05_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 { cuda_device::tcgen05::tcgen05_ld_32x32b_x32_raw(address + $offset) };)+ + let mut values = [0.0; $n]; + $( + let raw = $chunk.to_array(); + let mut i = 0; + while i < 32 { + cuda_device::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) }; + cuda_device::tcgen05::tcgen05_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 { + cuda_device::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) }; + cuda_device::tcgen05::tcgen05_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 { + cuda_device::thread::__unroll_config::<0>(); + chunk[i] = raw[$offset + i]; + i += 1; + } + unsafe { cuda_device::tcgen05::tcgen05_st_32x32b_x32_raw(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/cuteir/cute-rs/src/tensor.rs b/cuda-oxide/cuteir/cute-rs/src/tensor.rs index ae1ce711c7..6481a97625 100644 --- a/cuda-oxide/cuteir/cute-rs/src/tensor.rs +++ b/cuda-oxide/cuteir/cute-rs/src/tensor.rs @@ -47,6 +47,7 @@ impl TensorElement for f32 {} #[repr(C)] #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct Contiguous1D { + /// Number of logical elements in the contiguous row. len: usize, } @@ -87,6 +88,7 @@ impl Contiguous1D { #[repr(C)] #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct Zipped1D { + /// Full tensor length, used to retain a possibly partial final tile. len: usize, } @@ -107,7 +109,9 @@ impl Zipped1D { #[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, } @@ -142,7 +146,9 @@ impl Tile1D { /// The `N` values from one tile, held in thread-local registers. #[repr(transparent)] #[derive(Clone, Copy, Debug, PartialEq)] -pub struct RegisterTile([T; N]); +pub struct RegisterTile( + #[doc = "Per-thread register values in increasing tile-coordinate order."] [T; N], +); impl RegisterTile { /// Borrow the values as a fixed-size array. @@ -200,8 +206,11 @@ 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, } @@ -328,10 +337,15 @@ 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]>, } @@ -618,15 +632,18 @@ pub unsafe fn tensor_store_element_abs<'a, T: Copy + TensorElement, const TIL 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]; @@ -637,6 +654,7 @@ mod tests { 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]; @@ -655,6 +673,7 @@ mod tests { assert!(!tail.layout().is_full()); } + /// 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]); @@ -663,6 +682,7 @@ mod tests { 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!( diff --git a/cuda-oxide/cuteir/cute-rs/src/tile.rs b/cuda-oxide/cuteir/cute-rs/src/tile.rs index 8321838586..118a6f11d3 100644 --- a/cuda-oxide/cuteir/cute-rs/src/tile.rs +++ b/cuda-oxide/cuteir/cute-rs/src/tile.rs @@ -103,7 +103,9 @@ pub unsafe fn assume_div(x: usize) -> usize { // // 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 {} @@ -125,7 +127,9 @@ impl GlobalStoreElement for u32 {} /// 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(pub [T; N]); +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. /// diff --git a/cuda-oxide/cuteir/cute-rs/src/tiled_copy.rs b/cuda-oxide/cuteir/cute-rs/src/tiled_copy.rs index 856694d3bd..2c19de8b73 100644 --- a/cuda-oxide/cuteir/cute-rs/src/tiled_copy.rs +++ b/cuda-oxide/cuteir/cute-rs/src/tiled_copy.rs @@ -39,7 +39,9 @@ use crate::tensor::TensorElement; /// /// `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)>, } @@ -69,7 +71,9 @@ impl, T: Copy, LD, Role> GlobalCopyTensor, 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)>, } @@ -111,6 +115,7 @@ impl, T: Copy, Layout, Role> SharedTensor, T: Copy, Layout, Role>( base: *mut T, @@ -141,6 +146,7 @@ pub mod __compiler { /// The value contains no runtime data. The destination layout comes from the /// [`SharedTensor`] passed to [`Self::copy`]. pub struct TiledCopy { + /// Records the atom and thread/value/tile maps without runtime storage. config: PhantomData<(Atom, ThreadLayout, ValueLayout, TileLayout)>, } @@ -202,14 +208,18 @@ impl Default } } +/// Regression tests for this module’s layout, representation and rejection contracts. #[cfg(test)] mod tests { use super::*; + /// Copy plans must be zero-sized 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); + /// Distinct logical element marker used to test typed copy carrier size. #[derive(Clone, Copy)] struct Element; impl TensorElement for Element {} diff --git a/cuda-oxide/cuteir/cute-rs/src/tma.rs b/cuda-oxide/cuteir/cute-rs/src/tma.rs deleted file mode 100644 index bd45d451fa..0000000000 --- a/cuda-oxide/cuteir/cute-rs/src/tma.rs +++ /dev/null @@ -1,567 +0,0 @@ -/* - * 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: TmaDesc -> kernel: &TmaDesc -//! 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: &TmaDesc` 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 TmaDesc { - /// 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 TmaDesc { - /// Prefetch this descriptor into the device's tensor-map cache. - /// - /// 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(&self) { - unsafe { ctm_rs::tma::prefetch_descriptor(core::ptr::from_ref(self).cast()) }; - } -} - -// The type parameters describe the encoded bytes; they are not stored values. -impl Copy for TmaDesc {} - -impl Clone for TmaDesc { - 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 `make_tma_desc_2d`. 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 `make_tma_desc_2d_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::SWIZZLE_B, SmemL::SWIZZLE_M + log2, SmemL::SWIZZLE_S) -} - -/// 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), - } -} - -/// Number of bytes in one `SmemL` tile of `T`. -/// -/// 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 = tile_bytes::(); -/// ``` -#[inline(always)] -pub const fn tile_bytes() -> u32 { - (SmemL::ROWS * SmemL::COLS) 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::ROWS -/// col = c * SmemL::COLS -/// ``` -/// -/// The caller must initialize the `mbarrier`, attach [`tile_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 `make_tma_desc_2d` 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 TmaDesc, - coord: (usize, usize), - dst_smem_tile: &mut crate::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::ROWS, c * SmemL::COLS)`. -/// -/// 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::TmaStorePipeline`] -/// 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 TmaDesc, - coord: (usize, usize), - src_smem_tile: crate::SmemTile, -) { - let _ = (desc, coord, src_smem_tile); - unreachable!("cute-rs `copy_tma_s2g_2d` executed outside device compilation") -} - -#[cfg(feature = "host")] -mod host { - use super::{ - ReifySmem2D, TmaDesc, 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. - /// - /// 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 make_tma_desc_2d( - global_base: *mut core::ffi::c_void, - rows: u64, - cols: u64, - ld_elements: u64, - ) -> Result, alloc::string::String> { - make_tma_desc_2d_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 [`make_tma_desc_2d`]. 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 make_tma_desc_2d_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::ROWS <= 0 || SmemL::COLS <= 0 { - return Err(format!( - "TMA tile extents must be positive, got {} x {}", - SmemL::ROWS, - SmemL::COLS - )); - } - if SmemL::ROW_STRIDE != SmemL::COLS || SmemL::COL_STRIDE != 1 { - return Err(format!( - "TMA v0 requires a row-major shared tile, got strides ({}, {})", - SmemL::ROW_STRIDE, - SmemL::COL_STRIDE - )); - } - 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::SWIZZLE_B, - SmemL::SWIZZLE_M, - SmemL::SWIZZLE_S - )); - }; - if let Some(span_bytes) = tma_swizzle_span_bytes(mode) { - let pitch_bytes = SmemL::COLS - .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::COLS as u32, SmemL::ROWS 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() }; - const _: () = assert!(core::mem::size_of::() == 128); - Ok(TmaDesc { - bytes: unsafe { core::mem::transmute::(map) }, - marker: PhantomData, - }) - } -} - -#[cfg(feature = "host")] -pub use host::{TmaElement, make_tma_desc_2d, make_tma_desc_2d_with_options}; - -#[cfg(test)] -mod tests { - #[test] - fn descriptor_value_abi_does_not_require_copy_marker_types() { - // Layout and element markers need not be Copy: only descriptor bytes move. - struct Marker; - fn assert_copy() {} - assert_copy::>(); - assert_eq!(core::mem::size_of::>(), 128); - assert_eq!(core::mem::align_of::>(), 64); - } - - use super::*; - - #[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); - } - - #[test] - fn element_recast_reaches_the_byte_modes() { - use crate::markers::{Composed, ReifySmem2D, RowMajor, Swizzle}; - // For f16, element pattern S<3,3,3> becomes byte pattern S<3,4,3>, - // which is TMA B128. - type Smem = Composed, 0, RowMajor<8, 64>>; - assert_eq!(Smem::SWIZZLE_B, 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); - } - - #[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)); - } - - #[test] - fn flattened_mxfp4_ab_tile_is_8192_bytes_with_b64_swizzle() { - use crate::markers::{Composed, RowMajor, Swizzle}; - - // Two FP4 values share each byte: 128 K values become 64 bytes. - // Byte pattern S<2,4,3> is TMA B64. - type ByteSmem = Composed, 0, RowMajor<128, 64>>; - assert_eq!(tile_bytes::(), 8_192); - assert_eq!( - tma_swizzle_mode_for::(1), - Some(TmaSwizzleMode::B64) - ); - } - - #[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. - type ScaleSmem = RowMajor<1, 256>; - assert_eq!(tile_bytes::(), 512); - assert_eq!( - tma_swizzle_mode_for::(2), - Some(TmaSwizzleMode::None) - ); - } - - #[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); - } - - #[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/cuteir/cute-rs/src/tma/host.rs b/cuda-oxide/cuteir/cute-rs/src/tma/host.rs new file mode 100644 index 0000000000..bea5b76c12 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/tma/host.rs @@ -0,0 +1,186 @@ +/* + * 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, TmaDesc, 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. +/// +/// 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 make_tma_desc_2d( + global_base: *mut core::ffi::c_void, + rows: u64, + cols: u64, + ld_elements: u64, +) -> Result, alloc::string::String> { + make_tma_desc_2d_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 [`make_tma_desc_2d`]. 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 make_tma_desc_2d_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::ROWS <= 0 || SmemL::COLS <= 0 { + return Err(format!( + "TMA tile extents must be positive, got {} x {}", + SmemL::ROWS, + SmemL::COLS + )); + } + if SmemL::ROW_STRIDE != SmemL::COLS || SmemL::COL_STRIDE != 1 { + return Err(format!( + "TMA v0 requires a row-major shared tile, got strides ({}, {})", + SmemL::ROW_STRIDE, + SmemL::COL_STRIDE + )); + } + 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::SWIZZLE_B, + SmemL::SWIZZLE_M, + SmemL::SWIZZLE_S + )); + }; + if let Some(span_bytes) = tma_swizzle_span_bytes(mode) { + let pitch_bytes = SmemL::COLS + .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::COLS as u32, SmemL::ROWS 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(TmaDesc { + bytes: unsafe { core::mem::transmute::(map) }, + marker: PhantomData, + }) +} diff --git a/cuda-oxide/cuteir/cute-rs/src/tma/mod.rs b/cuda-oxide/cuteir/cute-rs/src/tma/mod.rs new file mode 100644 index 0000000000..cb04f4b735 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/tma/mod.rs @@ -0,0 +1,288 @@ +/* + * 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: TmaDesc -> kernel: &TmaDesc +//! 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: &TmaDesc` 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 TmaDesc { + /// 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 TmaDesc { + /// Prefetch this descriptor into the device's tensor-map cache. + /// + /// 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(&self) { + unsafe { ctm_rs::tma::prefetch_descriptor(core::ptr::from_ref(self).cast()) }; + } +} + +// The type parameters describe the encoded bytes; they are not stored values. +impl Copy for TmaDesc {} + +impl Clone for TmaDesc { + 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 `make_tma_desc_2d`. 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 `make_tma_desc_2d_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::SWIZZLE_B, SmemL::SWIZZLE_M + log2, SmemL::SWIZZLE_S) +} + +/// 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), + } +} + +/// Number of bytes in one `SmemL` tile of `T`. +/// +/// 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 = tile_bytes::(); +/// ``` +#[inline(always)] +pub const fn tile_bytes() -> u32 { + (SmemL::ROWS * SmemL::COLS) 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::ROWS +/// col = c * SmemL::COLS +/// ``` +/// +/// The caller must initialize the `mbarrier`, attach [`tile_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 `make_tma_desc_2d` 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 TmaDesc, + coord: (usize, usize), + dst_smem_tile: &mut crate::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::ROWS, c * SmemL::COLS)`. +/// +/// 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::TmaStorePipeline`] +/// 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 TmaDesc, + coord: (usize, usize), + src_smem_tile: crate::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, make_tma_desc_2d, make_tma_desc_2d_with_options}; + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cuteir/cute-rs/src/tma/tests.rs b/cuda-oxide/cuteir/cute-rs/src/tma/tests.rs new file mode 100644 index 0000000000..096ba269b9 --- /dev/null +++ b/cuda-oxide/cuteir/cute-rs/src/tma/tests.rs @@ -0,0 +1,114 @@ +/* + * 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::{Composed, 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 = Composed, 0, RowMajor<8, 64>>; + assert_eq!(Smem::SWIZZLE_B, 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::{Composed, 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 = Composed, 0, RowMajor<128, 64>>; + assert_eq!(tile_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!(tile_bytes::(), 512); + assert_eq!( + tma_swizzle_mode_for::(2), + Some(TmaSwizzleMode::None) + ); +} + +/// 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/cuteir/cute-rs/tests/arch.rs b/cuda-oxide/cuteir/cute-rs/tests/arch.rs index f39dd3a98d..10775248d5 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/arch.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/arch.rs @@ -1,6 +1,8 @@ // 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::assume_warp_uniform; diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_storage.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_storage.rs index 1d60cafe52..fd313cfed9 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_storage.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_storage.rs @@ -1,6 +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/cuteir/cute-rs/tests/compile_fail/global_store_storage.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_storage.stderr index 1f99fc57e3..b638b5b8f8 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_storage.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_storage.stderr @@ -1,7 +1,7 @@ error[E0277]: the trait bound `f64: GlobalStoreElement` is not satisfied - --> tests/compile_fail/global_store_storage.rs:5:38 + --> tests/compile_fail/global_store_storage.rs:7:38 | -5 | let _ = cute_rs::GlobalStoreTile([0.0f64; 4]); +7 | let _ = cute_rs::GlobalStoreTile([0.0f64; 4]); | ------------------------ ^^^^^^^^^^^ the trait `GlobalStoreElement` is not implemented for `f64` | | | required by a bound introduced by this call @@ -18,5 +18,5 @@ help: the following other types implement trait `GlobalStoreElement` note: required by a bound in `GlobalStoreTile` --> src/tile.rs | - | pub struct GlobalStoreTile(pub [T; N]); + | pub struct GlobalStoreTile( | ^^^^^^^^^^^^^^^^^^ required by this bound in `GlobalStoreTile` diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_unsafe.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_unsafe.rs index 5655b0b9f8..469e76e7cb 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_unsafe.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_unsafe.rs @@ -1,6 +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 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/cuteir/cute-rs/tests/compile_fail/global_store_unsafe.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_unsafe.stderr index 4d524fedc4..1b5e588dc7 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_unsafe.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/global_store_unsafe.stderr @@ -1,7 +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:6:5 + --> tests/compile_fail/global_store_unsafe.rs:8:5 | -6 | cute_rs::store_global_no_allocate(core::ptr::null_mut(), &values); +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/cuteir/cute-rs/tests/compile_fail/immutable_store.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/immutable_store.rs index 8e68c11056..763765ab0c 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/immutable_store.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/immutable_store.rs @@ -1,10 +1,13 @@ // 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) diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/immutable_store.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/immutable_store.stderr index f1ee65d854..882df9aa4f 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/immutable_store.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/immutable_store.stderr @@ -1,7 +1,7 @@ error[E0599]: no method named `store` found for struct `Tensor<'a, Element, Layout, Storage>` in the current scope - --> tests/compile_fail/immutable_store.rs:14:19 + --> tests/compile_fail/immutable_store.rs:17:19 | -14 | unsafe { tile.store(values) }; +17 | unsafe { tile.store(values) }; | ^^^^^ | help: there is a method `storage` with a similar name, but with different arguments diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/invalid_storage.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/invalid_storage.rs index 9fb201701a..edba62ad50 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/invalid_storage.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/invalid_storage.rs @@ -1,10 +1,13 @@ // 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::{E2M1, KMajor, Mkl, Tensor}; +/// Reject an FP32 carrier for packed E2M1 tensor storage. + fn main() { let wrong_carrier = [0.0f32; 32]; let _ = Tensor::::from_storage( diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/invalid_storage.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/invalid_storage.stderr index da4597c8cb..401a268d8e 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/invalid_storage.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/invalid_storage.stderr @@ -1,7 +1,7 @@ error[E0277]: the trait bound `E2M1: TensorElement` is not satisfied - --> tests/compile_fail/invalid_storage.rs:10:13 + --> tests/compile_fail/invalid_storage.rs:13:13 | -10 | let _ = Tensor::::from_storage( +13 | let _ = Tensor::::from_storage( | ^^^^^^^^^^^^^^^^^^^^^^ the trait `TensorElement` is not implemented for `E2M1` | help: the trait `TensorElement` is not implemented for `E2M1` @@ -20,9 +20,9 @@ note: required by a bound in `Tensor` | ^^^^^^^^^^^^^^^^^^^^^^ required by this bound in `Tensor` error[E0599]: the associated function or constant `from_storage` exists for struct `Tensor<'_, E2M1, _, f32>`, but its trait bounds were not satisfied - --> tests/compile_fail/invalid_storage.rs:10:37 + --> tests/compile_fail/invalid_storage.rs:13:37 | -10 | let _ = Tensor::::from_storage( +13 | let _ = Tensor::::from_storage( | ^^^^^^^^^^^^ associated function or constant cannot be called on `Tensor<'_, E2M1, _, f32>` due to unsatisfied trait bounds | ::: src/numeric.rs diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.rs index 592db4ddc3..cb90a4157f 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.rs @@ -1,12 +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 shared tile passed to the A-fragment loader. #![feature(f16)] use core::ptr; use cute_rs::{Mxf4E2M1, Mxfp4TiledMma, 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) diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.stderr index bbec330d8f..a6b0010806 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_operand_role_mismatch.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/mma_operand_role_mismatch.rs:15:33 + --> tests/compile_fail/mma_operand_role_mismatch.rs:18:33 | -15 | let _ = unsafe { mma.load_a(&shared_b, (0, 0)) }; +18 | let _ = unsafe { mma.load_a(&shared_b, (0, 0)) }; | ------ ^^^^^^^^^ expected `&SharedTensor`, found `&SharedTensor` | | | arguments to this method are incorrect @@ -9,7 +9,7 @@ error[E0308]: mismatched types = note: expected reference `&SharedTensor` found reference `&SharedTensor` note: method defined here - --> src/block_scaled_mma.rs + --> src/block_scaled_mma/tiled.rs | | pub unsafe fn load_a( | ^^^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.rs index d3e1f56456..4e30938f5f 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.rs @@ -1,10 +1,13 @@ // 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::{AccC, FragA, FragB, Mkl, Mxfp4BlockScaledTensor, Mxfp4TiledMma, 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]; diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.stderr index f457c4425f..2c2315dabc 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/mma_scale_role_mismatch.stderr @@ -1,27 +1,27 @@ error[E0308]: arguments to this method are incorrect - --> tests/compile_fail/mma_scale_role_mismatch.rs:16:13 + --> tests/compile_fail/mma_scale_role_mismatch.rs:19:13 | -16 | mma.gemm( +19 | mma.gemm( | ^^^^ -17 | FragA([0; 4]), -18 | vector.load_scale_pair(0, 0), +20 | FragA([0; 4]), +21 | vector.load_scale_pair(0, 0), | ---------------------------- expected `MmaScalePair`, found `MmaScalePair` -19 | FragB([0; 2]), -20 | matrix.load_scale_pair(0, 0), +22 | FragB([0; 2]), +23 | matrix.load_scale_pair(0, 0), | ---------------------------- expected `MmaScalePair`, found `MmaScalePair` | note: method defined here - --> src/block_scaled_mma.rs + --> src/block_scaled_mma/tiled.rs | | pub unsafe fn gemm( | ^^^^ help: swap these arguments | -16 ~ mma.gemm( -17 + FragA([0; 4]), -18 + matrix.load_scale_pair(0, 0), -19 + FragB([0; 2]), -20 + vector.load_scale_pair(0, 0), -21 + AccC::ZERO, -22 + ) +19 ~ mma.gemm( +20 + FragA([0; 4]), +21 + matrix.load_scale_pair(0, 0), +22 + FragB([0; 2]), +23 + vector.load_scale_pair(0, 0), +24 + AccC::ZERO, +25 + ) | diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.rs index 68ac5fc346..874d7a28a3 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.rs @@ -1,6 +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::sm100_attention::{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/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.stderr index aceb88a10d..12c09e1528 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.stderr @@ -1,26 +1,27 @@ error[E0080]: evaluation panicked: assertion failed: contiguous >= 32 && contiguous.is_multiple_of(32) - --> src/sm100_attention.rs + --> src/sm100_attention/shared.rs | | impl SharedLayout { + | /// Reject unsupported contiguous swizzle spans and outer tile extents. | const VALID: () = { ... | assert!(contiguous >= 32 && contiguous.is_multiple_of(32)); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::sm100_attention::SharedLayout::<128, 16>::VALID` failed here note: erroneous constant encountered - --> src/sm100_attention.rs + --> src/sm100_attention/shared.rs | | const VALID: () = SharedLayout::::VALID; | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ note: erroneous constant encountered - --> src/sm100_attention.rs + --> src/sm100_attention/mma.rs | | let () = AL::VALID; | ^^^^^^^^^ note: erroneous constant encountered - --> src/sm100_attention.rs + --> src/sm100_attention/mma.rs | | let () = Self::VALID; | ^^^^^^^^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_operand_role.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_operand_role.rs index 349fb8f34b..5da9e4748c 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_operand_role.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_operand_role.rs @@ -1,9 +1,12 @@ // 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::sm100_attention::{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()) }; diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_operand_role.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_operand_role.stderr index d12a1611a1..21d2ee5d81 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_operand_role.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_operand_role.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/sm100_attention_operand_role.rs:11:26 + --> tests/compile_fail/sm100_attention_operand_role.rs:14:26 | -11 | unsafe { mma.gemm(0, &a, &b, false) }; +14 | unsafe { mma.gemm(0, &a, &b, false) }; | ---- ^^ expected `OperandA`, found `OperandB` | | | arguments to this method are incorrect @@ -9,7 +9,7 @@ error[E0308]: mismatched types = note: expected reference `&SharedTile, OperandA>` found reference `&SharedTile, OperandB>` note: method defined here - --> src/sm100_attention.rs + --> src/sm100_attention/mma.rs | | pub unsafe fn gemm( | ^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.rs index 35b104bc50..411d70c88f 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.rs @@ -1,13 +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::sm100_attention::{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/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.stderr index 1a4ab1eeb6..af21e246e0 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.stderr @@ -1,5 +1,5 @@ error[E0080]: evaluation panicked: assertion failed: TILES > 0 - --> src/sm100_attention.rs + --> src/sm100_attention/shared.rs | | impl SharedTileArray { ... @@ -8,13 +8,13 @@ error[E0080]: evaluation panicked: assertion failed: TILES > 0 | ^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::sm100_attention::SharedTileArray::, cute_rs::OperandB, 0>::VALID` failed here note: erroneous constant encountered - --> src/sm100_attention.rs + --> src/sm100_attention/shared.rs | | let () = Self::VALID; | ^^^^^^^^^^^ error[E0080]: evaluation panicked: assertion failed: TILES <= u32::MAX / Self::TILE_BYTES - --> src/sm100_attention.rs + --> src/sm100_attention/shared.rs | | impl SharedTileArray { ... diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.rs index 2b65e4d3c3..7462d8bab8 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.rs @@ -1,9 +1,12 @@ // 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::sm100_attention::{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) }; diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.stderr index 4ecd0009b5..7c183b72de 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/sm100_attention_tile_array_role.rs:11:26 + --> tests/compile_fail/sm100_attention_tile_array_role.rs:14:26 | -11 | unsafe { mma.gemm(0, &tile, &tile, false) }; +14 | unsafe { mma.gemm(0, &tile, &tile, false) }; | ---- ^^^^^ expected `OperandA`, found `OperandB` | | | arguments to this method are incorrect @@ -9,7 +9,7 @@ error[E0308]: mismatched types = note: expected reference `&SharedTile, OperandA>` found reference `&SharedTile, OperandB>` note: method defined here - --> src/sm100_attention.rs + --> src/sm100_attention/mma.rs | | pub unsafe fn gemm( | ^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tma_layout.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tma_layout.rs index 00cf1f779f..1d614fc6c7 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tma_layout.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tma_layout.rs @@ -1,11 +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::sm100_attention::{SharedLayout, SharedTile, TmaCopy}; use cute_rs::{OperandB, TmaDesc}; +/// 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: &TmaDesc, 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/cuteir/cute-rs/tests/compile_fail/sm100_attention_tma_layout.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tma_layout.stderr index 7ac84cc134..5bc38c81b2 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tma_layout.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_tma_layout.stderr @@ -1,15 +1,15 @@ error[E0308]: mismatched types - --> tests/compile_fail/sm100_attention_tma_layout.rs:9:43 - | -9 | 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>` + --> 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/sm100_attention.rs - | - | pub unsafe fn load( - | ^^^^ + --> src/sm100_attention/tma.rs + | + | pub unsafe fn load( + | ^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.rs index 97d8eb31e1..451c0a19cb 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.rs @@ -1,7 +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::sm100_attention::{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/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.stderr index df70ab423b..5ec185c0dd 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_attention_ts_k16.stderr @@ -1,14 +1,15 @@ error[E0080]: evaluation panicked: assertion failed: K == 32 || (K > 0 && K.is_multiple_of(64)) - --> src/sm100_attention.rs + --> src/sm100_attention/mma.rs | | impl MmaTs { + | /// Check the group-one TMEM/shared profile, supported K span and B layout shape. | const VALID: () = { ... | assert!(K == 32 || (K > 0 && K.is_multiple_of(64))); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::sm100_attention::MmaTs::, 128, 128, 16, 1>::VALID` failed here note: erroneous constant encountered - --> src/sm100_attention.rs + --> src/sm100_attention/mma.rs | | let () = Self::VALID; | ^^^^^^^^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.rs index a084526bd7..bb574ad7db 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.rs @@ -1,9 +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::{Composed, RowMajor, Sm100TiledMma, Swizzle}; +/// Local 128-by-64 FP16 A tile with the B128 swizzle. type A = Composed, 0, RowMajor<128, 64>>; +/// Local 96-by-64 B tile belonging to the collective N192 partition. type B0 = Composed, 0, RowMajor<96, 64>>; +/// Local 80-by-64 B tile belonging to the collective N160 partition. type B1 = Composed, 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/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.stderr index 95495a17d5..7441774a5c 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.stderr @@ -1,5 +1,5 @@ error[E0080]: evaluation panicked: assertion failed: B0L::ROWS == (N0 / 2) as i64 && B0L::COLS == K as i64 - --> src/sm100.rs + --> src/sm100/mma.rs | | impl< ... @@ -9,7 +9,7 @@ error[E0080]: evaluation panicked: assertion failed: B0L::ROWS == (N0 / 2) as i6 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::Sm100TiledMma::, 0, cute_rs::RowMajor<128, 64>>, cute_rs::Composed, 0, cute_rs::RowMajor<80, 64>>, cute_rs::Composed, 0, cute_rs::RowMajor<96, 64>>, 192, 160>::VALID` failed here note: erroneous constant encountered - --> src/sm100.rs + --> src/sm100/mma.rs | | let () = Self::VALID; | ^^^^^^^^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.rs index f46deb54e0..918121a119 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.rs @@ -1,10 +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 supplied as the paired SM100 MMA A operand. #![feature(f16)] use cute_rs::{Composed, OperandB, RowMajor, Sm100SharedTile, Sm100TiledMma, Swizzle}; +/// Local 128-by-64 FP16 A tile with the B128 swizzle. type A = Composed, 0, RowMajor<128, 64>>; +/// Local 96-by-64 B tile belonging to the collective N192 partition. type B0 = Composed, 0, RowMajor<96, 64>>; +/// Local 80-by-64 B tile belonging to the collective N160 partition. type B1 = Composed, 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()) }; diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.stderr index 6256b3e68b..38dc037dd2 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/sm100_operand_role_mismatch.rs:14:21 + --> tests/compile_fail/sm100_operand_role_mismatch.rs:19:21 | -14 | mma.gemm(0, &a, &b0, &b1, false); +19 | mma.gemm(0, &a, &b0, &b1, false); | ---- ^^ expected `OperandA`, found `OperandB` | | | arguments to this method are incorrect @@ -9,7 +9,7 @@ error[E0308]: mismatched types = note: expected reference `&Sm100SharedTile, 0, RowMajor<128, 64>>, OperandA>` found reference `&Sm100SharedTile, 0, RowMajor<128, 64>>, OperandB>` note: method defined here - --> src/sm100.rs + --> src/sm100/mma.rs | | pub unsafe fn gemm( | ^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.rs index 1ab0f9c25c..e44d682c49 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.rs @@ -1,11 +1,15 @@ // 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::{ Composed, OperandA, RowMajor, Sm100ClusterTmaCopy, Sm100SharedTile, Swizzle, TmaDesc, }; +/// B128-swizzled 128-by-64 descriptor layout. type A = Composed, 0, RowMajor<128, 64>>; +/// B128-swizzled 96-by-64 destination layout. type B = Composed, 0, RowMajor<96, 64>>; +/// Reject a cluster TMA copy whose descriptor and destination shapes differ. fn main() { let desc: *const TmaDesc = core::ptr::null(); let mut dst = unsafe { Sm100SharedTile::::from_raw(core::ptr::null_mut()) }; diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.stderr index 7f5a32f6ad..8ef9675ba2 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/sm100_tma_layout_mismatch.rs:13:46 + --> tests/compile_fail/sm100_tma_layout_mismatch.rs:17:46 | -13 | Sm100ClusterTmaCopy::::copy(desc, &mut dst, 0, 0, core::ptr::null_mut(), 0); +17 | Sm100ClusterTmaCopy::::copy(desc, &mut dst, 0, 0, core::ptr::null_mut(), 0); | ------------------------------ ^^^^^^^^ expected `128`, found `96` | | | arguments to this function are incorrect @@ -9,7 +9,7 @@ error[E0308]: mismatched types = 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.rs + --> src/sm100/tma.rs | | pub unsafe fn copy( | ^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.rs index 7f37992682..bd158c04b3 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.rs @@ -1,12 +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/cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.stderr index 00586a86c4..6867d505d9 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_clone.stderr @@ -1,11 +1,11 @@ error[E0277]: the trait bound `TensorMut<'static, f32, Contiguous1D>: Clone` is not satisfied - --> tests/compile_fail/tensor_mut_clone.rs:11:21 + --> tests/compile_fail/tensor_mut_clone.rs:16:21 | -11 | require_clone::>(); +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:8:21 + --> tests/compile_fail/tensor_mut_clone.rs:11:21 | - 8 | fn require_clone() {} +11 | fn require_clone() {} | ^^^^^ required by this bound in `require_clone` diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.rs index cee2d69c45..2cdd40d0f1 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.rs @@ -1,12 +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/cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.stderr index 38f037f791..2d15b48954 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tensor_mut_copy.stderr @@ -1,11 +1,11 @@ error[E0277]: the trait bound `TensorMut<'static, f32, Contiguous1D>: Copy` is not satisfied - --> tests/compile_fail/tensor_mut_copy.rs:11:20 + --> tests/compile_fail/tensor_mut_copy.rs:16:20 | -11 | require_copy::>(); +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:8:20 + --> tests/compile_fail/tensor_mut_copy.rs:11:20 | - 8 | fn require_copy() {} +11 | fn require_copy() {} | ^^^^ required by this bound in `require_copy` diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.rs index c1dad21f29..9287d296cd 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.rs @@ -1,12 +1,15 @@ // 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::{CpAsync, Mkl, Mxfp4BlockScaledTensor, 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]; diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.stderr index e04950d934..2b46e75d78 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/tiled_copy_encoding_mismatch.rs:20:41 + --> tests/compile_fail/tiled_copy_encoding_mismatch.rs:23:41 | -20 | unsafe { copy.copy(&source, (0, 0), &mut wrong_encoding, 0) }; +23 | unsafe { copy.copy(&source, (0, 0), &mut wrong_encoding, 0) }; | ---- ^^^^^^^^^^^^^^^^^^^ expected `&mut SharedTensor`, found `&mut SharedTensor` | | | arguments to this method are incorrect diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.rs index 523c3a5544..fac1f68f97 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.rs @@ -1,6 +1,7 @@ // 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; @@ -9,6 +10,8 @@ use cute_rs::{ CpAsync, Mkl, Mxf4E2M1, Mxfp4BlockScaledTensor, 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]; diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr index 98337b9f4b..904bb4ebab 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/tiled_copy_role_mismatch.rs:22:41 + --> tests/compile_fail/tiled_copy_role_mismatch.rs:25:41 | -22 | unsafe { copy.copy(&source, (0, 0), &mut shared_b, 0) }; +25 | unsafe { copy.copy(&source, (0, 0), &mut shared_b, 0) }; | ---- ^^^^^^^^^^^^^ expected `&mut SharedTensor`, found `&mut SharedTensor<_, f16, (), Nkl>` | | | arguments to this method are incorrect diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs index 7012747c86..832f53beb6 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs @@ -3,10 +3,15 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Compile-fail probe: require an unsafe block for device-only descriptor prefetch. use cute_rs::TmaDesc; +/// Require an unsafe block for device-only descriptor prefetch. + fn prefetch(descriptor: &TmaDesc) { descriptor.prefetch(); } +/// Unused host entry point; the device or rejection probe is only type-checked. + fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.stderr index b0350be753..ee4764119f 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.stderr @@ -1,7 +1,7 @@ error[E0133]: call to unsafe function `TmaDesc::::prefetch` is unsafe and requires unsafe block - --> tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs:9:5 - | -9 | descriptor.prefetch(); - | ^^^^^^^^^^^^^^^^^^^^^ call to unsafe function - | - = note: consult the function's documentation for information on how to avoid undefined behavior + --> tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs:12:5 + | +12 | descriptor.prefetch(); + | ^^^^^^^^^^^^^^^^^^^^^ call to unsafe function + | + = note: consult the function's documentation for information on how to avoid undefined behavior diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.rs index 7e6c41b461..2a4f53f9d5 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.rs @@ -1,15 +1,21 @@ // 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, TmaDesc, 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:: { diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.stderr index 4ab478a8ff..beadbe5a27 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/tma_store_layout_mismatch.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/tma_store_layout_mismatch.rs:21:50 + --> tests/compile_fail/tma_store_layout_mismatch.rs:27:50 | -21 | unsafe { copy_tma_s2g_2d(descriptor, (0, 0), source) }; +27 | unsafe { copy_tma_s2g_2d(descriptor, (0, 0), source) }; | --------------- ^^^^^^ expected `8`, found `4` | | | arguments to this function are incorrect @@ -9,7 +9,7 @@ error[E0308]: mismatched types = note: expected struct `SmemTile>` found struct `SmemTile>` note: function defined here - --> src/tma.rs + --> src/tma/mod.rs | | pub unsafe fn copy_tma_s2g_2d( | ^^^^^^^^^^^^^^^ diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.rs index 888517576d..e90b2faa71 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.rs @@ -3,10 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Compile-fail probe: reject consumer cursors on producer operations and producer cursors on consumer operations. use cute_rs::sm100_attention::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()); @@ -15,4 +18,5 @@ unsafe fn wrong_roles(ring: TwoCtaRing<15, AsyncCompletion>) { 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/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.stderr index d7b48db9e5..363707ea9a 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:12:29 + --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:15:29 | -12 | ring.acquire_cursor(RingCursor::<15, Consumer>::new()); +15 | ring.acquire_cursor(RingCursor::<15, Consumer>::new()); | -------------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `RingCursor<15, Producer>`, found `RingCursor<15, Consumer>` | | | arguments to this method are incorrect @@ -15,9 +15,9 @@ note: method defined here | ^^^^^^^^^^^^^^ error[E0308]: mismatched types - --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:13:26 + --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:16:26 | -13 | ring.wait_cursor(RingCursor::<15, Producer>::new()); +16 | ring.wait_cursor(RingCursor::<15, Producer>::new()); | ----------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `RingCursor<15, Consumer>`, found `RingCursor<15, Producer>` | | | arguments to this method are incorrect @@ -31,9 +31,9 @@ note: method defined here | ^^^^^^^^^^^ error[E0308]: mismatched types - --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:14:30 + --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:17:30 | -14 | ring.mma_full_cursor(RingCursor::<15, Consumer>::new()); +17 | ring.mma_full_cursor(RingCursor::<15, Consumer>::new()); | --------------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `RingCursor<15, Producer>`, found `RingCursor<15, Consumer>` | | | arguments to this method are incorrect @@ -47,9 +47,9 @@ note: method defined here | ^^^^^^^^^^^^^^^ error[E0308]: mismatched types - --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:15:31 + --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:18:31 | -15 | ring.mma_empty_cursor(RingCursor::<15, Producer>::new()); +18 | ring.mma_empty_cursor(RingCursor::<15, Producer>::new()); | ---------------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `RingCursor<15, Consumer>`, found `RingCursor<15, Producer>` | | | arguments to this method are incorrect diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs index f7cdb0fa3a..3d6e235935 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs @@ -3,14 +3,18 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Compile-fail probe: reject cursors whose stage count differs from the ring. use cute_rs::sm100_attention::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/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.stderr index f08ac78904..413c15a524 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs:12:29 + --> tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs:15:29 | -12 | ring.acquire_cursor(RingCursor::<2, Producer>::new()); +15 | ring.acquire_cursor(RingCursor::<2, Producer>::new()); | -------------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `15`, found `2` | | | arguments to this method are incorrect @@ -15,9 +15,9 @@ note: method defined here | ^^^^^^^^^^^^^^ error[E0308]: mismatched types - --> tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs:13:26 + --> tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs:16:26 | -13 | ring.wait_cursor(RingCursor::<2, Consumer>::new()); +16 | ring.wait_cursor(RingCursor::<2, Consumer>::new()); | ----------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `15`, found `2` | | | arguments to this method are incorrect diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.rs index fb3cf1e50a..b1031029e3 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.rs @@ -3,7 +3,11 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Compile-fail probe: reject a zero-stage cursor during const evaluation. use cute_rs::sm100_attention::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/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.stderr index aa48eb812d..7d0dde218f 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.stderr @@ -2,6 +2,7 @@ error[E0080]: evaluation panicked: assertion failed: STAGES > 0 --> src/sm100_attention/pipeline.rs | | impl RingCursor { + | /// Reject zero-stage cursors before any slot arithmetic. | const VALID: () = assert!(STAGES > 0); | ^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::sm100_attention::pipeline::RingCursor::<0, cute_rs::sm100_attention::pipeline::Producer>::VALID` failed here diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs index d92b96f328..0666b96c79 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs @@ -3,10 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Compile-fail probe: reject completion signals unavailable for async-copy reuse and local-shared policies. use cute_rs::sm100_attention::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>, @@ -22,4 +25,6 @@ unsafe fn wrong_non_tc_signals( } } +/// Unused host entry point; the device or rejection probe is only type-checked. + fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.stderr index 93e2b2e79c..039e450e87 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.stderr @@ -1,7 +1,7 @@ error[E0277]: the trait bound `AsyncCopyReuse: SignalPolicy` is not satisfied - --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:16:16 + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:19:16 | -16 | copies.thread_empty(0); +19 | copies.thread_empty(0); | ^^^^^^^^^^^^ the trait `SignalPolicy` is not implemented for `AsyncCopyReuse` | help: the following other types implement trait `SignalPolicy` @@ -23,9 +23,9 @@ note: required by a bound in `TwoCtaRing::::thread_empty` | ^^^^^^^^^^^^ 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:17:15 + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:20:15 | -17 | reuse.expect(0, 16); +20 | reuse.expect(0, 16); | ^^^^^^ the trait `TmaPolicy` is not implemented for `AsyncCopyReuse` | help: the trait `TmaPolicy` is implemented for `AsyncCompletion` @@ -44,9 +44,9 @@ note: required by a bound in `TwoCtaRing::::expect` | ^^^^^^^^^ required by this bound in `TwoCtaRing::::expect` error[E0277]: the trait bound `LocalShared: TmaPolicy` is not satisfied - --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:18:15 + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:21:15 | -18 | pages.expect(0, 16); +21 | pages.expect(0, 16); | ^^^^^^ the trait `TmaPolicy` is not implemented for `LocalShared` | help: the trait `TmaPolicy` is implemented for `AsyncCompletion` @@ -65,9 +65,9 @@ note: required by a bound in `TwoCtaRing::::expect` | ^^^^^^^^^ required by this bound in `TwoCtaRing::::expect` error[E0277]: the trait bound `LocalShared: MmaPolicy` is not satisfied - --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:19:15 + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:22:15 | -19 | pages.mma_full(0); +22 | pages.mma_full(0); | ^^^^^^^^ the trait `MmaPolicy` is not implemented for `LocalShared` | help: the following other types implement trait `MmaPolicy` @@ -95,9 +95,9 @@ note: required by a bound in `TwoCtaRing::::mma_full` | ^^^^^^^^^ 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:20:15 + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:23:15 | -20 | pages.mma_empty_cursor(RingCursor::<4, Consumer>::new()); +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` @@ -125,9 +125,9 @@ note: required by a bound in `TwoCtaRing::::mma_empty_curso | ^^^^^^^^^ 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:21:15 + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:24:15 | -21 | pages.thread_full(0); +24 | pages.thread_full(0); | ^^^^^^^^^^^ the trait `SignalPolicy` is not implemented for `LocalShared` | help: the following other types implement trait `SignalPolicy` diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.rs index e672adac2d..d49d507012 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.rs @@ -3,9 +3,12 @@ * 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::sm100_attention::pipeline::WaitPolicy; +/// Invalid external wait policy with no ordering guarantees. + #[derive(Clone, Copy)] struct Unordered; @@ -13,4 +16,6 @@ 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/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.stderr index 20b3b6a35b..3b3609558c 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.stderr @@ -1,13 +1,13 @@ error[E0277]: the trait bound `Unordered: cute_rs::sm100_attention::pipeline::sealed::Wait` is not satisfied - --> tests/compile_fail/two_cta_ring_policy_sealed.rs:12:21 + --> tests/compile_fail/two_cta_ring_policy_sealed.rs:15:21 | -12 | impl WaitPolicy for Unordered { +15 | impl WaitPolicy for Unordered { | ^^^^^^^^^ unsatisfied trait bound | help: the trait `cute_rs::sm100_attention::pipeline::sealed::Wait` is not implemented for `Unordered` - --> tests/compile_fail/two_cta_ring_policy_sealed.rs:10:1 + --> tests/compile_fail/two_cta_ring_policy_sealed.rs:13:1 | -10 | struct Unordered; +13 | struct Unordered; | ^^^^^^^^^^^^^^^^ = help: the following other types implement trait `cute_rs::sm100_attention::pipeline::sealed::Wait`: AsyncCompletion diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.rs index a728b1a766..4392f20225 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.rs @@ -3,10 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Compile-fail probe: reject signaling mechanisms incompatible with each ring policy. use cute_rs::sm100_attention::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>, @@ -22,4 +25,6 @@ unsafe fn wrong_completion_mechanisms( } } +/// Unused host entry point; the device or rejection probe is only type-checked. + fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.stderr index ad7ded01d2..6a70aeff15 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.stderr @@ -1,7 +1,7 @@ error[E0277]: the trait bound `TensorMemory: TmaPolicy` is not satisfied - --> tests/compile_fail/two_cta_ring_wrong_signal.rs:17:16 + --> tests/compile_fail/two_cta_ring_wrong_signal.rs:20:16 | -17 | tensor.expect(0, 16); +20 | tensor.expect(0, 16); | ^^^^^^ the trait `TmaPolicy` is not implemented for `TensorMemory` | help: the trait `TmaPolicy` is implemented for `AsyncCompletion` @@ -20,9 +20,9 @@ note: required by a bound in `TwoCtaRing::::expect` | ^^^^^^^^^ required by this bound in `TwoCtaRing::::expect` error[E0277]: the trait bound `SharedToAsync: MmaPolicy` is not satisfied - --> tests/compile_fail/two_cta_ring_wrong_signal.rs:18:16 + --> tests/compile_fail/two_cta_ring_wrong_signal.rs:21:16 | -18 | shared.mma_full(0); +21 | shared.mma_full(0); | ^^^^^^^^ the trait `MmaPolicy` is not implemented for `SharedToAsync` | help: the following other types implement trait `MmaPolicy` @@ -50,9 +50,9 @@ note: required by a bound in `TwoCtaRing::::mma_full` | ^^^^^^^^^ 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:19:16 + --> tests/compile_fail/two_cta_ring_wrong_signal.rs:22:16 | -19 | shared.thread_empty(0); +22 | shared.thread_empty(0); | ^^^^^^^^^^^^ the trait `SignalPolicy` is not implemented for `SharedReuse` | help: the following other types implement trait `SignalPolicy` @@ -74,9 +74,9 @@ note: required by a bound in `TwoCtaRing::::thread_empty` | ^^^^^^^^^^^^ 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:20:16 + --> tests/compile_fail/two_cta_ring_wrong_signal.rs:23:16 | -20 | copies.thread_full(0); +23 | copies.thread_full(0); | ^^^^^^^^^^^ the trait `SignalPolicy` is not implemented for `AsyncCompletion` | help: the following other types implement trait `SignalPolicy` @@ -98,9 +98,9 @@ note: required by a bound in `TwoCtaRing::::thread_full` | ^^^^^^^^^^^^ 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:21:15 + --> tests/compile_fail/two_cta_ring_wrong_signal.rs:24:15 | -21 | local.mma_empty(0); +24 | local.mma_empty(0); | ^^^^^^^^^ the trait `MmaPolicy` is not implemented for `LocalMemory` | help: the following other types implement trait `MmaPolicy` diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.rs index 0e3a496eec..851d350aad 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.rs @@ -3,8 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Compile-fail probe: reject a zero-stage ring during const evaluation. use cute_rs::sm100_attention::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/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.stderr index 5ca941d8cd..51530555d4 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.stderr @@ -2,6 +2,7 @@ error[E0080]: evaluation panicked: assertion failed: STAGES > 0 --> src/sm100_attention/pipeline.rs | | impl TwoCtaRing { + | /// Reject zero-stage rings before modulo or storage-size calculations. | const VALID: () = assert!(STAGES > 0); | ^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::sm100_attention::pipeline::TwoCtaRing::<0, cute_rs::sm100_attention::pipeline::AsyncCompletion>::VALID` failed here diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.rs index ab9ead80e5..3ad1332471 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.rs @@ -1,6 +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::assume_warp_uniform(0); } diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.stderr index cd6bc430fe..a8c8dc27b6 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_unsafe.stderr @@ -1,7 +1,7 @@ error[E0133]: call to unsafe function `assume_warp_uniform` is unsafe and requires unsafe block - --> tests/compile_fail/warp_uniform_unsafe.rs:5:13 + --> tests/compile_fail/warp_uniform_unsafe.rs:7:13 | -5 | let _ = cute_rs::arch::assume_warp_uniform(0); +7 | let _ = cute_rs::arch::assume_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/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.rs index bb435abb77..ef3747ff58 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.rs @@ -1,6 +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::assume_warp_uniform(0u64) }; } diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.stderr b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.stderr index 929edb2930..01eaf0762d 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.stderr +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_fail/warp_uniform_width.stderr @@ -1,7 +1,7 @@ error[E0308]: mismatched types - --> tests/compile_fail/warp_uniform_width.rs:5:57 + --> tests/compile_fail/warp_uniform_width.rs:7:57 | -5 | let _ = unsafe { cute_rs::arch::assume_warp_uniform(0u64) }; +7 | let _ = unsafe { cute_rs::arch::assume_warp_uniform(0u64) }; | ---------------------------------- ^^^^ expected `u32`, found `u64` | | | arguments to this function are incorrect @@ -13,6 +13,6 @@ note: function defined here | ^^^^^^^^^^^^^^^^^^^ help: change the type of the numeric literal from `u64` to `u32` | -5 - let _ = unsafe { cute_rs::arch::assume_warp_uniform(0u64) }; -5 + let _ = unsafe { cute_rs::arch::assume_warp_uniform(0u32) }; +7 - let _ = unsafe { cute_rs::arch::assume_warp_uniform(0u64) }; +7 + let _ = unsafe { cute_rs::arch::assume_warp_uniform(0u32) }; | diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_pass/tma_descriptor_prefetch.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/tma_descriptor_prefetch.rs index 40fb4e7cb5..abd3215904 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_pass/tma_descriptor_prefetch.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/tma_descriptor_prefetch.rs @@ -3,15 +3,20 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Compile-pass probe: type-check descriptor prefetch for arbitrary element and layout markers. use cute_rs::TmaDesc; // 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: &TmaDesc) { unsafe { descriptor.prefetch() }; } +/// 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(&TmaDesc) = prefetch; let _: unsafe fn(&TmaDesc) = prefetch; diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_non_tc_receivers.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_non_tc_receivers.rs index 971eac0bad..36a0ac33f9 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_non_tc_receivers.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_non_tc_receivers.rs @@ -3,10 +3,13 @@ * 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::sm100_attention::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>, @@ -33,4 +36,6 @@ unsafe fn non_tc_receivers( } } +/// Unused host entry point; the device or rejection probe is only type-checked. + fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_cursor.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_cursor.rs index fc6a4dd959..d4a1bc6f8d 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_cursor.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_cursor.rs @@ -3,11 +3,14 @@ * 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::sm100_attention::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>, @@ -29,4 +32,5 @@ unsafe fn valid_cursors( 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/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_signal.rs b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_signal.rs index d80f9114ff..b579e72510 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_signal.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/compile_pass/two_cta_ring_valid_signal.rs @@ -3,10 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Compile-pass probe: type-check the completion signals permitted by each ring policy. use cute_rs::sm100_attention::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>, @@ -25,4 +28,6 @@ unsafe fn valid_completion_mechanisms( } } +/// Unused host entry point; the device or rejection probe is only type-checked. + fn main() {} diff --git a/cuda-oxide/cuteir/cute-rs/tests/global_store.rs b/cuda-oxide/cuteir/cute-rs/tests/global_store.rs index 5dd4bbb5bc..ea4a1aae27 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/global_store.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/global_store.rs @@ -1,6 +1,8 @@ // 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) = @@ -9,6 +11,7 @@ fn global_store_preserves_element_and_vector_width_in_its_signature() { 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}; @@ -20,6 +23,7 @@ fn global_store_carrier_guarantees_the_full_source_alignment() { 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(); diff --git a/cuda-oxide/cuteir/cute-rs/tests/sm100_attention.rs b/cuda-oxide/cuteir/cute-rs/tests/sm100_attention.rs index 7f72d6f067..d0dee730ed 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/sm100_attention.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/sm100_attention.rs @@ -3,9 +3,11 @@ * SPDX-License-Identifier: Apache-2.0 */ +#![doc = "Checks attention shared layouts, typed MMA shapes and complete tile-array contracts."] use cute_rs::sm100_attention::{MmaSs, MmaTs, SharedLayout, SharedTileArray, SmemLayout}; use cute_rs::{OperandA, OperandB}; +/// 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 { @@ -20,6 +22,7 @@ fn check_bijection() { ); } +/// Supported K-major and MN-major tiled swizzles must cover each element exactly once. #[test] fn tiled_swizzles_cover_storage_without_overlap() { check_bijection::>(); @@ -29,8 +32,10 @@ fn tiled_swizzles_cover_storage_without_overlap() { 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::offset(0, 0), 0); assert_eq!(K::offset(0, 64), 8192); @@ -47,6 +52,7 @@ fn tma_contiguous_chunks_and_transposed_v_share_the_same_storage() { 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(); @@ -56,9 +62,13 @@ fn independent_qk_and_pv_configurations_are_constructible() { } // 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([u16; 3 * 16384]); +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(); @@ -95,6 +105,7 @@ fn check_tile_array() { 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>(); @@ -106,6 +117,7 @@ fn complete_tile_arrays_preserve_addresses_and_swizzles() { 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(); @@ -115,6 +127,7 @@ fn tile_arrays_reject_operand_role_changes_and_invalid_extents() { // 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]); diff --git a/cuda-oxide/cuteir/cute-rs/tests/sm100_attention_pipeline.rs b/cuda-oxide/cuteir/cute-rs/tests/sm100_attention_pipeline.rs index 63f2c62a27..dbea966019 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/sm100_attention_pipeline.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/sm100_attention_pipeline.rs @@ -3,12 +3,14 @@ * 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::sm100_attention::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. @@ -25,6 +27,7 @@ fn barrier_views_select_disjoint_arrays_and_wrap_slots() { } } +/// 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(); @@ -36,6 +39,7 @@ fn completion_policies_reject_incompatible_signals() { 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]; @@ -65,6 +69,7 @@ fn cursors_preserve_phase_and_addresses_across_persistent_work() { 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(); @@ -77,6 +82,7 @@ fn one_stage_cursors_flip_opposite_phases_on_every_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(); diff --git a/cuda-oxide/cuteir/cute-rs/tests/type_safety.rs b/cuda-oxide/cuteir/cute-rs/tests/type_safety.rs index eba6a93fb2..fb2b987ee7 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/type_safety.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/type_safety.rs @@ -3,6 +3,8 @@ * 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(); diff --git a/cuda-oxide/cuteir/dialect-cute/src/attributes.rs b/cuda-oxide/cuteir/dialect-cute/src/attributes.rs deleted file mode 100644 index d30344cd20..0000000000 --- a/cuda-oxide/cuteir/dialect-cute/src/attributes.rs +++ /dev/null @@ -1,1092 +0,0 @@ -/* - * 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. - -use core::fmt; - -use combine::Parser; -use cute_layout::{ComposedLayout, Layout, OffsetUnit, ParseLayoutError, Swizzle}; -use pliron::common_traits::Verify; -use pliron::context::Context; -use pliron::impl_printable_for_display; -use pliron::irfmt::parsers::quoted_string_parser; -use pliron::parsable::{Parsable, ParseResult, StateStream}; -use pliron::result::Error; -use pliron::verify_err_noloc; -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(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(pub u32); - -/// 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 { - Producer, - 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 { - pub role: CutePipelineRoleAttr, - 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(()) - } -} - -/// 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 { - pub m_tiles: u64, - pub n_tiles: u64, - 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(()) - } -} - -/// 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 { - 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 { - ReadOnly, - 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 { - Plain, - E2M1, - UE8M0, -} - -/// 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 { - Generic, - Mkl, - 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 { - Contiguous1D, - Zipped1D(u64), - Tile1D(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(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(()) - } -} - -/// 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 { - Full, - Row, - KTile(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(()) - } -} - -/// 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(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(()) - } -} - -/// 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, - B, - BPair, - BPairKMajor, -} - -/// A complete swizzled layout, stored as one typed attribute. -/// -/// The quoted payload is intentionally simple and round-trippable: -/// -/// ```text -/// "3,4,3;0;elements;(8,32):(32,1)" -/// B M S off unit inner layout -/// ``` -#[pliron_attr(name = "cute.composed_layout", verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub struct CuteComposedLayoutAttr(pub ComposedLayout); - -impl fmt::Display for CuteComposedLayoutAttr { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let swizzle = self.0.outer(); - write!( - f, - "\"{},{},{};{};{};{}\"", - swizzle.bits, - swizzle.base, - swizzle.shift, - self.0.offset(), - self.0.unit(), - self.0.inner() - ) - } -} - -impl_printable_for_display!(CuteComposedLayoutAttr); - -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 inner = 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 outer = Swizzle::try_new(bits, base, shift).ok_or(ParseLayoutError)?; - ComposedLayout::new(outer, offset, inner, 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() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use cute_layout::IntTuple; - - #[test] - fn composed_layout_text_round_trips_every_addressing_piece() { - let inner = 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, inner, 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); - } - - #[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()); - } - - #[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.outer(), Swizzle::new(2, 0, -3)); - assert_eq!(attr.to_string(), "\"2,0,-3;0;bytes;16:4\""); - } - - #[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()); - } -} - -/// 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(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() - } -} - -/// 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 { - 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 { - ScaleStage, - ScaleK64, - A, - 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 { - pub m: u64, - pub n: u64, - pub k: u64, - pub a_format: CuteTensorFormatAttr, - pub b_format: CuteTensorFormatAttr, - pub scale_format: CuteTensorFormatAttr, - pub values_per_scale: u64, - pub accumulator: CuteMmaAccumulatorAttr, - pub threads: u32, - pub a_registers: u32, - pub b_registers: u32, - pub accumulator_registers: u32, -} - -impl CuteMmaAtomAttr { - /// SM120's packed MXFP4 `m16n8k64` atom. - #[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, - } - } -} - -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 - || self.scale_format != CuteTensorFormatAttr::UE8M0 - { - return verify_err_noloc!("cute.mma_atom v0 needs E2M1 A/B values with UE8M0 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 { - pub atom: CuteMmaAtomAttr, - pub cta_m: u64, - pub cta_n: u64, - pub cta_k: u64, - pub warp_m: u32, - pub warp_n: u32, - pub b_load_group: u32, - pub shared_layout: CuteComposedLayoutAttr, - pub m_ownership: CuteLayoutAttr, - 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"), - ), - } - } - - /// 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)) - } -} - -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.inner().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.leaves().first().copied() != Some(i64::from(self.warp_m)) - || self.n_ownership.0.shape.leaves().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(()) - } -} - -/// 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 { - pub stages: u32, -} - -impl CuteTmaStorePipelineAttr { - #[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 { - pub barrier_id: u32, - pub first_warp: u32, - pub warp_count: u32, - pub cta_warps: u32, - pub lanes_per_warp: u32, -} - -impl CuteCountedCtaBarrierAttr { - #[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(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 { - pub tiled_mma: CuteTiledMmaPlanAttr, - pub half_layout: CuteComposedLayoutAttr, - pub halves: u32, - 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 inner: Layout = "(128,64):(64,1)" - .parse() - .expect("fixed epilogue half layout is valid"); - let half_layout = - ComposedLayout::new(Swizzle::new(3, 3, 3), 0, inner, OffsetUnit::Elements) - .expect("fixed epilogue swizzle is valid"); - Self { - tiled_mma, - half_layout: CuteComposedLayoutAttr(half_layout), - halves: 2, - base_alignment: CuteAlignmentAttr(1024), - } - } - - #[must_use] - pub fn half_elements(&self) -> Option { - u64::try_from(self.half_layout.0.inner().checked_size()?).ok() - } - - #[must_use] - pub fn full_elements(&self) -> Option { - self.half_elements()?.checked_mul(u64::from(self.halves)) - } - - #[must_use] - pub fn half_bytes(&self, element_bytes: u64) -> Option { - self.half_elements()?.checked_mul(element_bytes) - } - - #[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_inner: 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_inner, 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.inner().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/cuteir/dialect-cute/src/attributes/copy.rs b/cuda-oxide/cuteir/dialect-cute/src/attributes/copy.rs new file mode 100644 index 0000000000..c12e525d7d --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/dialect-cute/src/attributes/epilogue.rs b/cuda-oxide/cuteir/dialect-cute/src/attributes/epilogue.rs new file mode 100644 index 0000000000..36bbe7bf1a --- /dev/null +++ b/cuda-oxide/cuteir/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 inner: Layout = "(128,64):(64,1)" + .parse() + .expect("fixed epilogue half layout is valid"); + let half_layout = + ComposedLayout::new(Swizzle::new(3, 3, 3), 0, inner, 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.inner().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_inner: 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_inner, 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.inner().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/cuteir/dialect-cute/src/attributes/gemv.rs b/cuda-oxide/cuteir/dialect-cute/src/attributes/gemv.rs new file mode 100644 index 0000000000..15df89d103 --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/dialect-cute/src/attributes/layout.rs b/cuda-oxide/cuteir/dialect-cute/src/attributes/layout.rs new file mode 100644 index 0000000000..1b74f5312c --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/attributes/layout.rs @@ -0,0 +1,122 @@ +/* + * 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: +/// +/// ```text +/// "3,4,3;0;elements;(8,32):(32,1)" +/// B M S off unit inner layout +/// ``` +#[pliron_attr(name = "cute.composed_layout", verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteComposedLayoutAttr( + #[doc = "Structural swizzle, offset, and inner layout."] pub ComposedLayout, +); + +impl fmt::Display for CuteComposedLayoutAttr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let swizzle = self.0.outer(); + write!( + f, + "\"{},{},{};{};{};{}\"", + swizzle.bits, + swizzle.base, + swizzle.shift, + self.0.offset(), + self.0.unit(), + self.0.inner() + ) + } +} + +impl_printable_for_display!(CuteComposedLayoutAttr); + +/// Parse the structural swizzle, offset unit, and inner 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 inner = 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 outer = Swizzle::try_new(bits, base, shift).ok_or(ParseLayoutError)?; + ComposedLayout::new(outer, offset, inner, 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/cuteir/dialect-cute/src/attributes/mma.rs b/cuda-oxide/cuteir/dialect-cute/src/attributes/mma.rs new file mode 100644 index 0000000000..61f4a131e1 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/attributes/mma.rs @@ -0,0 +1,353 @@ +/* + * 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. + #[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, + } + } +} + +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 + || self.scale_format != CuteTensorFormatAttr::UE8M0 + { + return verify_err_noloc!("cute.mma_atom v0 needs E2M1 A/B values with UE8M0 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"), + ), + } + } + + /// 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.inner().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.leaves().first().copied() != Some(i64::from(self.warp_m)) + || self.n_ownership.0.shape.leaves().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/cuteir/dialect-cute/src/attributes/mod.rs b/cuda-oxide/cuteir/dialect-cute/src/attributes/mod.rs new file mode 100644 index 0000000000..522967231a --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/dialect-cute/src/attributes/pipeline.rs b/cuda-oxide/cuteir/dialect-cute/src/attributes/pipeline.rs new file mode 100644 index 0000000000..1207147f8b --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/dialect-cute/src/attributes/scheduler.rs b/cuda-oxide/cuteir/dialect-cute/src/attributes/scheduler.rs new file mode 100644 index 0000000000..e7d706c237 --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/dialect-cute/src/attributes/tensor.rs b/cuda-oxide/cuteir/dialect-cute/src/attributes/tensor.rs new file mode 100644 index 0000000000..f66583f4a2 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/attributes/tensor.rs @@ -0,0 +1,187 @@ +/* + * 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, +} + +/// 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/cuteir/dialect-cute/src/attributes/tests.rs b/cuda-oxide/cuteir/dialect-cute/src/attributes/tests.rs new file mode 100644 index 0000000000..6d296f5ec7 --- /dev/null +++ b/cuda-oxide/cuteir/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 inner = 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, inner, 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.outer(), 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/cuteir/dialect-cute/src/copy_ops.rs b/cuda-oxide/cuteir/dialect-cute/src/copy_ops.rs index 730f79965f..91047e89bb 100644 --- a/cuda-oxide/cuteir/dialect-cute/src/copy_ops.rs +++ b/cuda-oxide/cuteir/dialect-cute/src/copy_ops.rs @@ -24,6 +24,7 @@ use pliron_derive::{pliron_attr, pliron_op}; #[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 { @@ -61,6 +62,7 @@ pub fn global_store_element_bytes(ctx: &Context, element: TypeHandle) -> Option< #[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( @@ -75,9 +77,11 @@ impl CuteStoreGlobalNoAllocateOp { 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) } diff --git a/cuda-oxide/cuteir/dialect-cute/src/epilogue_ops.rs b/cuda-oxide/cuteir/dialect-cute/src/epilogue_ops.rs index 191cbf7991..2c01ef1d85 100644 --- a/cuda-oxide/cuteir/dialect-cute/src/epilogue_ops.rs +++ b/cuda-oxide/cuteir/dialect-cute/src/epilogue_ops.rs @@ -49,6 +49,7 @@ use crate::attributes::{ 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) @@ -57,6 +58,7 @@ fn is_u32(ctx: &Context, value: Value) -> bool { .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) @@ -65,12 +67,14 @@ fn is_u64(ctx: &Context, value: Value) -> bool { .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) @@ -79,6 +83,7 @@ fn tma_view_of(ctx: &Context, value: Value) -> Option { .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) @@ -91,6 +96,7 @@ fn shared_pointer_matches(ctx: &Context, value: Value, storage: TypeHandle) -> b }) } +/// Require a valid epilogue tile attribute before checking carrier operands. fn checked_tile( tile: Option, ctx: &Context, @@ -101,6 +107,7 @@ fn checked_tile( Ok(tile) } +/// Build the supported 128-by-64 swizzled FP16 half-tile placement. fn current_half_layout() -> ComposedLayout { let inner: Layout = "(128,64):(64,1)" .parse() @@ -119,6 +126,7 @@ fn current_half_layout() -> ComposedLayout { 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( @@ -134,20 +142,24 @@ impl CuteEpilogueSmemOverlayOp { 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)?; @@ -197,6 +209,7 @@ impl Verify for CuteEpilogueSmemOverlayOp { 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, @@ -222,40 +235,48 @@ impl CuteEpilogueWarpSliceOp { 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)?; @@ -307,6 +328,7 @@ impl Verify for CuteEpilogueWarpSliceOp { 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, @@ -329,30 +351,36 @@ impl CuteEpilogueStoreFragmentOp { 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)?; @@ -442,27 +470,32 @@ impl CuteEpilogueSyncOp { 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) @@ -529,6 +562,7 @@ impl Verify for CuteEpilogueSyncOp { 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, @@ -551,31 +585,37 @@ impl CuteEpilogueHalfOp { 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) @@ -615,6 +655,7 @@ impl Verify for CuteEpilogueHalfOp { } } +/// Require a positive store-ring stage count on a zero-result, zero-operand operation. fn check_store_pipeline( pipeline: Option, ctx: &Context, @@ -639,6 +680,7 @@ fn check_store_pipeline( 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), @@ -647,10 +689,12 @@ impl CuteTmaStoreAcquireOp { 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) @@ -684,6 +728,7 @@ impl Verify for CuteTmaStoreAcquireOp { 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), @@ -692,10 +737,12 @@ impl CuteTmaStoreCommitOp { 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) @@ -729,6 +776,7 @@ impl Verify for CuteTmaStoreCommitOp { 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), @@ -737,10 +785,12 @@ impl CuteTmaStoreTailOp { 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) @@ -777,6 +827,7 @@ impl Verify for CuteTmaStoreTailOp { 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, @@ -796,25 +847,30 @@ impl CuteTmaStore2dSemanticOp { } } + /// 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) @@ -883,6 +939,7 @@ impl Verify for CuteTmaStore2dSemanticOp { } } +/// Regression fixtures covering accepted contracts and rejected malformed IR. #[cfg(test)] mod tests { use super::*; @@ -893,6 +950,7 @@ mod tests { 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() @@ -900,24 +958,37 @@ mod tests { .get_result(0) } + /// Build the canonical shared placement consumed by the output-tile fixture. fn data_layout() -> ComposedLayout { let inner: Layout = "(128,32):(32,1)".parse().unwrap(); ComposedLayout::new(Swizzle::new(2, 3, 3), 0, inner, 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); @@ -969,6 +1040,7 @@ mod tests { ) } + /// 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(); @@ -1036,6 +1108,7 @@ mod tests { ); } + /// 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(); @@ -1062,6 +1135,7 @@ mod tests { ); } + /// 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(); @@ -1109,6 +1183,7 @@ mod tests { 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(); @@ -1143,6 +1218,7 @@ mod tests { 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(); diff --git a/cuda-oxide/cuteir/dialect-cute/src/gemm_tma_ops.rs b/cuda-oxide/cuteir/dialect-cute/src/gemm_tma_ops.rs index 5d2f3e8450..61d0e67e6a 100644 --- a/cuda-oxide/cuteir/dialect-cute/src/gemm_tma_ops.rs +++ b/cuda-oxide/cuteir/dialect-cute/src/gemm_tma_ops.rs @@ -52,6 +52,7 @@ use crate::attributes::{ 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) @@ -60,6 +61,7 @@ fn tma_view_of(ctx: &Context, value: Value) -> Option { .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() { @@ -77,6 +79,7 @@ fn transport_element_bytes(ctx: &Context, element: TypeHandle) -> Option { } } +/// Check that the SSA value has an unsigned 64-bit integer type. fn is_u64(ctx: &Context, value: Value) -> bool { value .get_type(ctx) @@ -85,6 +88,7 @@ fn is_u64(ctx: &Context, value: Value) -> bool { .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, @@ -124,6 +128,7 @@ fn transport_view_type( 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, @@ -188,6 +193,7 @@ impl CuteTmaGmemViewOp { ) } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -251,6 +257,7 @@ impl Verify for CuteTmaGmemViewOp { 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, @@ -280,6 +287,7 @@ impl CuteTmaSmemViewOp { } } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -369,6 +377,7 @@ impl Verify for CuteTmaSmemViewOp { 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, @@ -395,30 +404,36 @@ impl CuteTmaCopy2dOp { } } + /// 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) @@ -491,6 +506,7 @@ impl Verify for CuteTmaCopy2dOp { } } +/// Regression fixtures covering accepted contracts and rejected malformed IR. #[cfg(test)] mod tests { use super::*; @@ -498,6 +514,7 @@ mod tests { 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() @@ -505,30 +522,42 @@ mod tests { .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 inner: Layout = "(128,64):(64,1)".parse().unwrap(); ComposedLayout::new(Swizzle::new(2, 4, 3), 0, inner, OffsetUnit::Elements).unwrap() } + /// Build the canonical scale transport tile used by the four-copy stage fixture. fn scale_layout() -> ComposedLayout { let inner: Layout = "(1,256):(256,1)".parse().unwrap(); ComposedLayout::new(Swizzle::IDENTITY, 0, inner, 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); @@ -562,6 +591,7 @@ mod tests { ) } + /// Build a detached transport copy from locally verified source and destination views. fn copy( ctx: &mut Context, descriptor: Value, @@ -580,6 +610,7 @@ mod tests { 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(); @@ -621,6 +652,7 @@ mod tests { } } + /// 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(); @@ -638,6 +670,7 @@ mod tests { 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(); @@ -676,6 +709,7 @@ mod tests { ); } + /// 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(); @@ -700,6 +734,7 @@ mod tests { ); } + /// 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(); @@ -721,6 +756,7 @@ mod tests { ); } + /// Descriptor-backed transport preserves the immutable tensor-map pointer contract. #[test] fn global_view_rejects_a_mutable_descriptor_pointer() { let (mut ctx, inputs) = setup(); @@ -739,6 +775,7 @@ mod tests { ); } + /// 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(); diff --git a/cuda-oxide/cuteir/dialect-cute/src/gemv_ops.rs b/cuda-oxide/cuteir/dialect-cute/src/gemv_ops.rs index 44187b411f..fc4a81fb2b 100644 --- a/cuda-oxide/cuteir/dialect-cute/src/gemv_ops.rs +++ b/cuda-oxide/cuteir/dialect-cute/src/gemv_ops.rs @@ -46,6 +46,7 @@ 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) @@ -54,6 +55,7 @@ fn tensor_view_of(ctx: &Context, value: Value) -> Option { .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) @@ -62,6 +64,7 @@ fn scaled_view_of(ctx: &Context, value: Value) -> Option { .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) @@ -70,6 +73,7 @@ fn fragment_of(ctx: &Context, value: Value) -> Option { .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) @@ -78,14 +82,17 @@ fn is_u64(ctx: &Context, value: Value) -> bool { .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 @@ -93,6 +100,7 @@ fn same_tensor_facts_except_role(left: &CuteTensorViewType, right: &CuteTensorVi && 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; @@ -128,6 +136,7 @@ fn fragments_can_dot(ctx: &Context, matrix: &CuteFragmentType, vector: &CuteFrag && 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() { @@ -249,6 +258,7 @@ impl CuteTensorMake2DOp { ) } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -347,6 +357,7 @@ impl Verify for CuteTensorMake2DOp { 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"); @@ -370,6 +381,7 @@ impl CuteScaledViewMakeOp { } } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -440,6 +452,7 @@ impl Verify for CuteScaledViewMakeOp { 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"); @@ -456,6 +469,7 @@ impl CuteScaledViewRowOp { } } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -509,6 +523,7 @@ impl Verify for CuteScaledViewRowOp { 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"); @@ -527,6 +542,7 @@ impl CuteScaledViewKTileOp { } } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -606,6 +622,7 @@ impl Verify for CuteScaledViewKTileOp { 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, @@ -628,6 +645,7 @@ impl CuteScaledViewLoadOp { load } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -716,6 +734,7 @@ impl Verify for CuteScaledViewLoadOp { 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( @@ -729,6 +748,7 @@ impl CuteDotOp { } } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -765,12 +785,14 @@ impl Verify for CuteDotOp { } } +/// Regression fixtures covering accepted contracts and rejected malformed IR. #[cfg(test)] mod tests { 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() @@ -778,15 +800,23 @@ mod tests { .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); @@ -815,10 +845,12 @@ mod tests { ) } + /// 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, @@ -837,6 +869,7 @@ mod tests { (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, @@ -853,6 +886,7 @@ mod tests { (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(); @@ -886,6 +920,7 @@ mod tests { ); } + /// 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(); @@ -921,6 +956,7 @@ mod tests { ); } + /// 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(); @@ -950,6 +986,7 @@ mod tests { ); } + /// 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(); @@ -985,6 +1022,7 @@ mod tests { ); } + /// 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(); @@ -1022,6 +1060,7 @@ mod tests { ); } + /// 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(); @@ -1059,6 +1098,7 @@ mod tests { 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(); diff --git a/cuda-oxide/cuteir/dialect-cute/src/lib.rs b/cuda-oxide/cuteir/dialect-cute/src/lib.rs index c0970fbac7..187f729eaf 100644 --- a/cuda-oxide/cuteir/dialect-cute/src/lib.rs +++ b/cuda-oxide/cuteir/dialect-cute/src/lib.rs @@ -43,8 +43,10 @@ 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, diff --git a/cuda-oxide/cuteir/dialect-cute/src/ops.rs b/cuda-oxide/cuteir/dialect-cute/src/ops.rs deleted file mode 100644 index 59ef970cfb..0000000000 --- a/cuda-oxide/cuteir/dialect-cute/src/ops.rs +++ /dev/null @@ -1,1609 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! cute dialect ops. -//! -//! The op set is intentionally small. `cute.copy` and `cute.copy_g2s` are -//! zero-result memory ops; `cute.assume_div` is an identity value that -//! preserves one compiler fact. - -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, OneOpdInterface, OneResultInterface}, - types::IntegerType, - }, - common_traits::Verify, - context::{Context, Ptr}, - location::Located, - 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, - CuteMatrixRoleAttr, -}; - -/// 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 { - #[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 - } - - pub fn wrap(op: Ptr) -> Self { - Self { op } - } -} - -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(()) - } -} - -/// Preserve the promise `value % divisor == 0` in SSA form. -/// -/// ```text -/// Rust cute IR lowering -/// ───── ─────── ──────── -/// y = assume_div::<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 { - 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 - } - - pub fn wrap(op: Ptr) -> Self { - Self { op } - } - - 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(()) - } -} - -/// 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(()) - } -} - -/// 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 { - 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 - } - - 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.inner().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(()) - } -} - -/// 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 { - 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 - } - - 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 { - 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 - } - - 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(()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use cute_layout::{IntTuple, MAX_STATIC_VALIDATION_BYTES, OffsetUnit, Swizzle}; - use dialect_mir::ops::MirUndefOp; - use dialect_mir::types::MirArrayType; - use pliron::builtin::types::{FP32Type, Signedness}; - - fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { - MirUndefOp::new(ctx, ty) - .get_operation() - .deref(ctx) - .get_result(0) - } - - 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) - } - - #[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(); - } - } - } - - #[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}" - ); - } - - 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) - } - - #[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, - ) - } - - #[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, - ) - } - - 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}`" - ); - } - - 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) - } - - 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) - } - - #[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}" - ); - } - - #[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(); - } - - #[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}"); - } - - #[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}"); - } - - #[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"); - } - - #[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"); - } - - #[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", - ); - } - - #[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", - ); - } - - #[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"); - } - - #[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"); - } - - #[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"); - } - - #[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/cuteir/dialect-cute/src/ops/copy.rs b/cuda-oxide/cuteir/dialect-cute/src/ops/copy.rs new file mode 100644 index 0000000000..a2a6475185 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/ops/copy.rs @@ -0,0 +1,382 @@ +/* + * 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 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}, + location::Located, + 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/cuteir/dialect-cute/src/ops/divisibility.rs b/cuda-oxide/cuteir/dialect-cute/src/ops/divisibility.rs new file mode 100644 index 0000000000..d2eb398158 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/ops/divisibility.rs @@ -0,0 +1,114 @@ +/* + * 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}, + location::Located, + 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_div::<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/cuteir/dialect-cute/src/ops/ldmatrix.rs b/cuda-oxide/cuteir/dialect-cute/src/ops/ldmatrix.rs new file mode 100644 index 0000000000..5df22c28d1 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/ops/ldmatrix.rs @@ -0,0 +1,170 @@ +/* + * 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}, + location::Located, + 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.inner().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/cuteir/dialect-cute/src/ops/mod.rs b/cuda-oxide/cuteir/dialect-cute/src/ops/mod.rs new file mode 100644 index 0000000000..3de3067b27 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/ops/mod.rs @@ -0,0 +1,28 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! cute dialect ops. +//! +//! The op set is intentionally small. `cute.copy` and `cute.copy_g2s` are +//! zero-result memory ops; `cute.assume_div` is an identity value that +//! preserves one compiler fact. + +/// Typed copy widths, divisibility promises, and warp matrix operand roles. +mod copy; +/// SSA identity operations preserving explicit integer divisibility promises. +mod divisibility; +/// Warp-cooperative matrix loads with typed operand roles and shared placement. +mod ldmatrix; +/// Descriptor-backed raw transport views and their direct copy contracts. +mod tma; + +pub use copy::{CuteCopyG2SOp, CuteCopyOp}; +pub use divisibility::CuteAssumeDivOp; +pub use ldmatrix::CuteLdmatrixOp; +pub use tma::{CuteTmaLoad2dOp, CuteTmaStore2dOp}; + +/// Regression fixtures covering accepted contracts and rejected malformed IR. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cuteir/dialect-cute/src/ops/tests.rs b/cuda-oxide/cuteir/dialect-cute/src/ops/tests.rs new file mode 100644 index 0000000000..341e9ea0e7 --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/dialect-cute/src/ops/tma.rs b/cuda-oxide/cuteir/dialect-cute/src/ops/tma.rs new file mode 100644 index 0000000000..364344b94c --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/ops/tma.rs @@ -0,0 +1,277 @@ +/* + * 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}, + location::Located, + op::Op, + operation::Operation, + result::Error, + r#type::{TypeHandle, Typed}, + value::Value, + verify_err, +}; +use pliron_derive::pliron_op; + +use super::copy::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/cuteir/dialect-cute/src/pipeline_ops.rs b/cuda-oxide/cuteir/dialect-cute/src/pipeline_ops.rs index e817f3158d..75e2d5a585 100644 --- a/cuda-oxide/cuteir/dialect-cute/src/pipeline_ops.rs +++ b/cuda-oxide/cuteir/dialect-cute/src/pipeline_ops.rs @@ -48,16 +48,20 @@ 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) @@ -66,6 +70,7 @@ fn is_u32(ctx: &Context, value: Value) -> bool { .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) @@ -74,6 +79,7 @@ fn is_u64(ctx: &Context, value: Value) -> bool { .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) @@ -82,6 +88,7 @@ fn pipeline_type_of(ctx: &Context, value: Value) -> Option Option { let producer = value.defining_op()?; if Operation::get_opid(producer, ctx) != CuteTmaLoadPipelineMakeOp::get_opid_static() { @@ -91,6 +98,7 @@ fn pipeline_make_of(ctx: &Context, value: Value) -> Option Option { value .get_type(ctx) @@ -139,6 +149,7 @@ fn barrier_pointer_of(ctx: &Context, value: Value) -> Option { .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; @@ -153,6 +164,7 @@ fn valid_shared_barrier_pointer(ctx: &Context, value: Value) -> bool { .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); [ @@ -210,6 +222,7 @@ impl CuteTmaLoadPipelineMakeOp { operation } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -299,6 +312,7 @@ impl Verify for CuteTmaLoadPipelineMakeOp { 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( @@ -312,10 +326,12 @@ impl CuteTmaLoadPipelineInitOp { } } + /// 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) @@ -360,6 +376,7 @@ impl Verify for CuteTmaLoadPipelineInitOp { 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 { @@ -376,20 +393,24 @@ impl CutePipelineStateNewOp { 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) @@ -429,6 +450,7 @@ impl Verify for CutePipelineStateNewOp { 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( @@ -444,10 +466,12 @@ impl CutePipelineStateSlotOp { 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) @@ -459,6 +483,7 @@ impl CutePipelineStateSlotOp { 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) @@ -498,6 +523,7 @@ impl Verify for CutePipelineStateSlotOp { 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 { @@ -514,30 +540,36 @@ impl CutePipelineStateAdvanceOp { 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) @@ -578,18 +610,22 @@ impl Verify for CutePipelineStateAdvanceOp { } } +/// 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) @@ -607,6 +643,7 @@ macro_rules! pipeline_state_accessors { 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, @@ -628,12 +665,14 @@ impl CutePipelineProducerAcquireOp { 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) @@ -686,6 +725,7 @@ impl Verify for CutePipelineProducerAcquireOp { 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, @@ -709,6 +749,7 @@ impl CutePipelineProducerExpectTxOp { operation } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -774,6 +815,7 @@ impl Verify for CutePipelineProducerExpectTxOp { 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, @@ -795,12 +837,14 @@ impl CutePipelineConsumerWaitOp { 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) @@ -850,6 +894,7 @@ impl Verify for CutePipelineConsumerWaitOp { pub struct CutePipelineConsumerReleaseOp; impl CutePipelineConsumerReleaseOp { + /// Build a detached empty-stage arrival after consumer reads finish. pub fn new( ctx: &mut Context, pipeline: Value, @@ -870,6 +915,7 @@ impl CutePipelineConsumerReleaseOp { operation } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -920,6 +966,7 @@ impl Verify for CutePipelineConsumerReleaseOp { 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, @@ -941,12 +988,14 @@ impl CutePipelineProducerTailOp { 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) @@ -986,15 +1035,20 @@ impl Verify for CutePipelineProducerTailOp { } } +/// Regression fixtures covering accepted contracts and rejected malformed IR. #[cfg(test)] mod tests { 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() @@ -1002,19 +1056,26 @@ mod tests { .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); @@ -1033,6 +1094,7 @@ mod tests { } } + /// Creation, initialization, cursor updates, and producer/consumer effects retain their separate operation boundaries. #[test] fn load_pipeline_stays_a_composable_ring_protocol() { let Setup { @@ -1124,6 +1186,7 @@ mod tests { ); } + /// Zero/overflowing stages and unsupported transaction or participation counts fail configuration checks. #[test] fn pipeline_config_rejects_impossible_rings() { let ctx = Context::new(); @@ -1164,6 +1227,7 @@ mod tests { 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 { @@ -1244,6 +1308,7 @@ mod tests { 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 { @@ -1307,6 +1372,7 @@ mod tests { ); } + /// 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 { @@ -1348,6 +1414,7 @@ mod tests { ); } + /// 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(); diff --git a/cuda-oxide/cuteir/dialect-cute/src/scheduler_ops.rs b/cuda-oxide/cuteir/dialect-cute/src/scheduler_ops.rs index b897dafc65..fda8c872b1 100644 --- a/cuda-oxide/cuteir/dialect-cute/src/scheduler_ops.rs +++ b/cuda-oxide/cuteir/dialect-cute/src/scheduler_ops.rs @@ -39,14 +39,17 @@ 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) @@ -55,6 +58,7 @@ fn is_u64(ctx: &Context, value: Value) -> bool { .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) @@ -63,6 +67,7 @@ fn is_bool(ctx: &Context, value: Value) -> bool { .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) @@ -85,6 +90,7 @@ fn work_tile_of(ctx: &Context, value: Value) -> Option { 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( @@ -100,6 +106,7 @@ impl CuteSchedulerHasWorkOp { operation } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -160,6 +167,7 @@ impl Verify for CuteSchedulerHasWorkOp { 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 { @@ -176,6 +184,7 @@ impl CuteSchedulerNew1dOp { operation } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -234,6 +243,7 @@ impl Verify for CuteSchedulerNew1dOp { 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 { @@ -250,6 +260,7 @@ impl CuteSchedulerCurrentOp { operation } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -317,6 +328,7 @@ impl Verify for CuteSchedulerCurrentOp { 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 { @@ -331,30 +343,36 @@ impl CuteWorkTileCoordinatesOp { } } + /// 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) @@ -408,6 +426,7 @@ impl Verify for CuteWorkTileCoordinatesOp { 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( @@ -423,20 +442,24 @@ impl CuteSchedulerAdvanceOp { 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) @@ -474,15 +497,18 @@ impl Verify for CuteSchedulerAdvanceOp { } } +/// Regression fixtures covering accepted contracts and rejected malformed IR. #[cfg(test)] mod tests { 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() @@ -490,6 +516,7 @@ mod tests { .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(); @@ -517,6 +544,7 @@ mod tests { 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(); @@ -537,6 +565,7 @@ mod tests { ); } + /// 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(); @@ -555,6 +584,7 @@ mod tests { ); } + /// 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(); @@ -564,6 +594,7 @@ mod tests { 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(); diff --git a/cuda-oxide/cuteir/dialect-cute/src/sm100_attention_ops.rs b/cuda-oxide/cuteir/dialect-cute/src/sm100_attention_ops.rs index 588e85083f..a699595cb5 100644 --- a/cuda-oxide/cuteir/dialect-cute/src/sm100_attention_ops.rs +++ b/cuda-oxide/cuteir/dialect-cute/src/sm100_attention_ops.rs @@ -14,6 +14,12 @@ //! 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}, @@ -35,6 +41,25 @@ use pliron::{ 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::sm100_attention_ops::{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 @@ -72,19 +97,35 @@ pub fn shared_layout(rows: u32, cols: u32, mn: bool) -> Result bool { let ty = value.get_type(ctx).deref(ctx); match want { @@ -198,12 +255,19 @@ fn carrier(ctx: &Context, value: Value, want: Carrier) -> bool { } } } +/// Generate independent zero-result MMA/TMA operations with typed plans and carrier checks. macro_rules! operation{ - ($ty:ident,$name:literal,$key:ident,$attr:ty,[$($arg:ident),*])=>{ + ($(#[$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{ @@ -218,6 +282,12 @@ macro_rules! operation{ } } 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, @@ -225,6 +295,12 @@ operation!( [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, @@ -232,6 +308,12 @@ operation!( [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, @@ -239,12 +321,19 @@ operation!( [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!( @@ -255,18 +344,24 @@ macro_rules! each_op { ) }; } +/// 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 diff --git a/cuda-oxide/cuteir/dialect-cute/src/sm100_ops.rs b/cuda-oxide/cuteir/dialect-cute/src/sm100_ops.rs index 134a7d3fc2..86d3d39472 100644 --- a/cuda-oxide/cuteir/dialect-cute/src/sm100_ops.rs +++ b/cuda-oxide/cuteir/dialect-cute/src/sm100_ops.rs @@ -14,6 +14,11 @@ //! 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}; @@ -37,16 +42,23 @@ use crate::layout::{ComposedLayout, Layout, OffsetUnit, Swizzle}; #[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 { @@ -63,24 +75,42 @@ impl Verify for CuteSm100TmemPlanAttr { } } +/// 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, } +/// 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 inner: Layout = format!("({rows},{columns}):({columns},1)") .parse() @@ -140,14 +170,22 @@ impl Verify for CuteSm100MmaPlanAttr { } } +/// 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 { @@ -171,14 +209,21 @@ impl Verify for CuteSm100ClusterCopyPlanAttr { } } +/// 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 { @@ -191,17 +236,46 @@ impl Verify for CuteSm100EpiloguePlanAttr { } } +/// 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::sm100_ops::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 { @@ -234,17 +308,27 @@ impl Verify for CuteSm100PipelinePlanAttr { } } +/// 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); @@ -293,21 +377,29 @@ fn carrier_matches(ctx: &Context, value: Value, carrier: Carrier) -> bool { } } +/// Generate zero-result operations with a typed plan and local carrier verification. macro_rules! semantic_op { - ($ty:ident, $name:literal, $key:ident, $plan:ty, [$($carrier:ident),*], $check:expr) => { + ($(#[$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 { @@ -331,6 +423,11 @@ macro_rules! semantic_op { }; } 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, @@ -339,6 +436,10 @@ semantic_op!( |_| 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, @@ -347,6 +448,12 @@ semantic_op!( |_| 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, @@ -355,6 +462,12 @@ semantic_op!( |_| 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, @@ -363,6 +476,12 @@ semantic_op!( |_| 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, @@ -371,6 +490,10 @@ semantic_op!( |_| 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, @@ -379,6 +502,10 @@ semantic_op!( |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, @@ -387,6 +514,10 @@ semantic_op!( |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, @@ -395,6 +526,10 @@ semantic_op!( |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, @@ -403,6 +538,10 @@ semantic_op!( |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, @@ -411,6 +550,10 @@ semantic_op!( |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, @@ -419,6 +562,10 @@ semantic_op!( |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, @@ -427,6 +574,9 @@ semantic_op!( |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, @@ -435,6 +585,10 @@ semantic_op!( |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, @@ -443,6 +597,10 @@ semantic_op!( |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, @@ -451,6 +609,11 @@ semantic_op!( |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, @@ -459,6 +622,9 @@ semantic_op!( |_| 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, @@ -467,6 +633,10 @@ semantic_op!( |_| 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, @@ -475,6 +645,9 @@ semantic_op!( |_| 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, @@ -483,6 +656,7 @@ semantic_op!( |_| true ); +/// Apply a callback to the closed paired-SM100 operation set. macro_rules! for_all_ops { ($callback:ident) => { $callback!( @@ -510,11 +684,13 @@ macro_rules! for_all_ops { } /// 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); @@ -527,11 +703,13 @@ pub fn register(ctx: &mut Context) { /// 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 } +/// Find and clone the operation attribute of the requested static plan type. fn static_plan(ctx: &Context, op: Ptr) -> Option { op.deref(ctx) .attributes @@ -543,6 +721,7 @@ fn static_plan(ctx: &Context, op: Ptr) -> Optio // 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()?; @@ -721,6 +900,7 @@ pub fn verify_story(ctx: &Context, ops: &[Ptr]) -> Result<(), String> Ok(()) } +/// Regression fixtures covering accepted contracts and rejected malformed IR. #[cfg(test)] mod tests { use super::*; @@ -736,17 +916,23 @@ mod tests { 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( @@ -765,6 +951,7 @@ mod tests { 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, @@ -776,6 +963,7 @@ mod tests { 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>, @@ -820,6 +1008,7 @@ mod tests { } } + /// Build the supported paired M256/N192+160/K64 plan with canonical shared layouts. fn mma_plan() -> CuteSm100MmaPlanAttr { CuteSm100MmaPlanAttr { m: 256, @@ -834,6 +1023,7 @@ mod tests { 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 }, @@ -844,11 +1034,13 @@ mod tests { 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); @@ -883,6 +1075,7 @@ mod tests { 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); @@ -942,6 +1135,7 @@ mod tests { (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(); @@ -958,6 +1152,7 @@ mod tests { 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(); @@ -984,12 +1179,14 @@ mod tests { 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(); @@ -998,6 +1195,7 @@ mod tests { 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(); @@ -1006,6 +1204,7 @@ mod tests { 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(); @@ -1014,6 +1213,7 @@ mod tests { 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(); @@ -1022,6 +1222,7 @@ mod tests { 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(); @@ -1030,6 +1231,7 @@ mod tests { 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(); @@ -1042,6 +1244,7 @@ mod tests { "{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(); @@ -1055,6 +1258,7 @@ mod tests { .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(); @@ -1068,6 +1272,7 @@ mod tests { .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(); diff --git a/cuda-oxide/cuteir/dialect-cute/src/smem_mma_ops.rs b/cuda-oxide/cuteir/dialect-cute/src/smem_mma_ops.rs index e1fc3fa1ef..b8731aec59 100644 --- a/cuda-oxide/cuteir/dialect-cute/src/smem_mma_ops.rs +++ b/cuda-oxide/cuteir/dialect-cute/src/smem_mma_ops.rs @@ -60,6 +60,7 @@ use crate::attributes::{ 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) @@ -68,6 +69,7 @@ fn is_u32(ctx: &Context, value: Value) -> bool { .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) @@ -76,6 +78,7 @@ fn is_u64(ctx: &Context, value: Value) -> bool { .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) @@ -84,10 +87,12 @@ fn is_f32(ctx: &Context, value: Value) -> bool { .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) @@ -100,6 +105,7 @@ fn shared_pointer_matches(ctx: &Context, value: Value, storage: TypeHandle) -> b }) } +/// Check the shared pointer and unsigned capacity carrier against the static tensor view. fn shared_carriers_match_view( ctx: &Context, base: Value, @@ -111,6 +117,7 @@ fn shared_carriers_match_view( }) } +/// Build the supported dense placement for one stage of block-scale carriers. fn canonical_scale_placement() -> ComposedLayout { let inner: Layout = "(32,4):(4,1)" .parse() @@ -119,6 +126,7 @@ fn canonical_scale_placement() -> ComposedLayout { .expect("identity scale placement is valid") } +/// Build the supported swizzled placement for one stage of packed FP4 values. fn canonical_data_placement() -> ComposedLayout { let inner: Layout = "(128,32):(32,1)" .parse() @@ -127,12 +135,14 @@ fn canonical_data_placement() -> ComposedLayout { .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")?; @@ -154,6 +164,7 @@ fn check_plan(plan: &CuteTiledMmaPlanAttr, ctx: &Context) -> Result<(), &'static Ok(()) } +/// Require the packed-value format, operand role, and shared placement needed by tiled MMA. fn check_data_view( ctx: &Context, view: &CuteSmemTensorType, @@ -181,6 +192,7 @@ fn check_data_view( Ok(()) } +/// Require the scale format, operand role, and canonical stage placement needed by tiled MMA. fn check_scale_view( ctx: &Context, view: &CuteSmemTensorType, @@ -209,6 +221,7 @@ fn check_scale_view( Ok(()) } +/// Require the ordinary aggregate's register counts to match its semantic fragment kind. fn check_carrier( ctx: &Context, value: Value, @@ -236,6 +249,7 @@ fn check_carrier( 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( @@ -251,25 +265,30 @@ impl CuteSmemTensorOverlayOp { 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) @@ -336,6 +355,7 @@ impl Verify for CuteSmemTensorOverlayOp { 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( @@ -351,20 +371,24 @@ impl CuteTiledMmaSliceOp { 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)) @@ -413,6 +437,7 @@ impl Verify for CuteTiledMmaSliceOp { 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, @@ -435,25 +460,30 @@ impl CuteFragmentFillOp { 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) @@ -505,6 +535,7 @@ impl Verify for CuteFragmentFillOp { 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, @@ -544,10 +575,12 @@ impl CuteMmaLoadScalesOp { 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)) @@ -567,6 +600,7 @@ impl CuteMmaLoadScalesOp { 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) @@ -644,6 +678,7 @@ impl Verify for CuteMmaLoadScalesOp { 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, @@ -665,15 +700,18 @@ impl CuteFragmentSliceKOp { 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) @@ -728,6 +766,7 @@ impl Verify for CuteFragmentSliceKOp { 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, @@ -755,10 +794,12 @@ impl CuteMmaLoadAOp { 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)) @@ -771,6 +812,7 @@ impl CuteMmaLoadAOp { 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) @@ -840,6 +882,7 @@ impl Verify for CuteMmaLoadAOp { 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, @@ -869,10 +912,12 @@ impl CuteMmaPartitionBOp { 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)) @@ -885,21 +930,25 @@ impl CuteMmaPartitionBOp { 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) @@ -971,6 +1020,7 @@ impl Verify for CuteMmaPartitionBOp { 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, @@ -1009,10 +1059,12 @@ impl CuteTiledGemmOp { 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)) @@ -1025,6 +1077,7 @@ impl CuteTiledGemmOp { 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) @@ -1098,6 +1151,7 @@ impl Verify for CuteTiledGemmOp { } } +/// Regression fixtures covering accepted contracts and rejected malformed IR. #[cfg(test)] mod tests { use super::*; @@ -1106,6 +1160,7 @@ mod tests { 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() @@ -1113,19 +1168,27 @@ mod tests { .get_result(0) } + /// Build the canonical swizzled packed-data shared layout for the MMA fixture. fn data_placement() -> ComposedLayout { let inner: Layout = "(128,32):(32,1)".parse().unwrap(); ComposedLayout::new(Swizzle::new(2, 3, 3), 0, inner, 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, @@ -1147,24 +1210,41 @@ mod tests { 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); @@ -1256,6 +1336,7 @@ mod tests { ) } + /// Attach semantic shared-tensor metadata to fixture pointer/capacity carriers. fn overlay( ctx: &mut Context, base_ty: TypeHandle, @@ -1267,6 +1348,7 @@ mod tests { 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(); @@ -1404,6 +1486,7 @@ mod tests { ); } + /// 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(); @@ -1436,6 +1519,7 @@ mod tests { ); } + /// 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(); @@ -1479,6 +1563,7 @@ mod tests { 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(); @@ -1523,6 +1608,7 @@ mod tests { ); } + /// 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(); diff --git a/cuda-oxide/cuteir/dialect-cute/src/tensor_ops.rs b/cuda-oxide/cuteir/dialect-cute/src/tensor_ops.rs index 9ac2290dd0..b8c1f86200 100644 --- a/cuda-oxide/cuteir/dialect-cute/src/tensor_ops.rs +++ b/cuda-oxide/cuteir/dialect-cute/src/tensor_ops.rs @@ -42,11 +42,13 @@ use crate::attributes::{ }; 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, @@ -57,6 +59,7 @@ fn transformed_view_type( 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) @@ -64,12 +67,14 @@ fn is_u64(ctx: &Context, value: Value) -> bool { .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 @@ -80,6 +85,7 @@ fn same_view_facts(left: &CuteTensorViewType, right: &CuteTensorViewType) -> boo && 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); @@ -91,6 +97,7 @@ fn carrier_pointer(ctx: &Context, value: Value, pointee: TypeHandle, needs_write && [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, @@ -159,6 +166,7 @@ impl CuteTensorMakeOp { } } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -238,6 +246,7 @@ impl Verify for CuteTensorMakeOp { 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); @@ -255,6 +264,7 @@ impl CuteTensorZippedDivideOp { divide } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -331,6 +341,7 @@ impl Verify for CuteTensorZippedDivideOp { 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"); @@ -353,6 +364,7 @@ impl CuteTensorSliceOp { } } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -414,6 +426,7 @@ impl Verify for CuteTensorSliceOp { 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 { @@ -428,6 +441,7 @@ impl CuteTensorIsFullOp { } } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -469,6 +483,7 @@ impl Verify for CuteTensorIsFullOp { 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 { @@ -483,6 +498,7 @@ impl CuteTensorBaseOp { } } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -531,6 +547,7 @@ impl Verify for CuteTensorBaseOp { 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, @@ -551,6 +568,7 @@ impl CuteTensorLoadIntoOp { load } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -640,6 +658,7 @@ impl CuteTensorStoreFromOp { store } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -711,6 +730,7 @@ impl Verify for CuteTensorStoreFromOp { 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( @@ -724,6 +744,7 @@ impl CuteTensorStoreElementAbsOp { } } + /// View an existing operation through this wrapper without allocating or verifying it. pub fn wrap(op: Ptr) -> Self { Self { op } } @@ -772,12 +793,14 @@ impl Verify for CuteTensorStoreElementAbsOp { } } +/// Regression fixtures covering accepted contracts and rejected malformed IR. #[cfg(test)] mod tests { 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() @@ -785,6 +808,7 @@ mod tests { .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); @@ -804,6 +828,7 @@ mod tests { (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(); @@ -870,6 +895,7 @@ mod tests { ); } + /// 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(); @@ -907,6 +933,7 @@ mod tests { ); } + /// 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(); diff --git a/cuda-oxide/cuteir/dialect-cute/src/types.rs b/cuda-oxide/cuteir/dialect-cute/src/types.rs deleted file mode 100644 index 94f5e946e5..0000000000 --- a/cuda-oxide/cuteir/dialect-cute/src/types.rs +++ /dev/null @@ -1,1220 +0,0 @@ -/* - * 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. - -use dialect_mir::types::{MirArrayType, MirFP16Type, MirStructType, MirTupleType}; -use pliron::builtin::types::{FP32Type, IntegerType}; -use pliron::common_traits::Verify; -use pliron::context::Context; -use pliron::derive::pliron_type; -use pliron::location::Location; -use pliron::result::Error; -use pliron::r#type::{Type, TypeHandle, TypedHandle}; -use pliron::verify_err; - -use crate::attributes::{ - CuteAlignmentAttr, CuteComposedLayoutAttr, CuteEpiloguePlanAttr, CuteScaledLayoutAttr, - CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, CuteTensorFormatAttr, CuteTensorLayoutAttr, - CuteTensorRoleAttr, CuteTileGridAttr, -}; - -/// 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 { - pub u32_registers: u64, - pub f32_registers: u64, -} - -impl MmaCarrierShape { - #[must_use] - pub const fn u32(registers: u64) -> Self { - Self { - u32_registers: registers, - f32_registers: 0, - } - } - - #[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) - } - - 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)?, - }) - } - - 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)?) - }) -} - -/// 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 { - pub logical: TypeHandle, - pub storage: TypeHandle, - pub format: CuteTensorFormatAttr, - pub role: CuteTensorRoleAttr, - pub space: CuteTensorAddressSpaceAttr, - pub access: CuteTensorAccessAttr, - pub alignment: CuteAlignmentAttr, - 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, - ) - } - - /// Size of one v0 storage element in bytes. - #[must_use] - pub fn storage_bytes(&self, ctx: &Context) -> Option { - let storage = self.storage.deref(ctx); - if storage.downcast_ref::().is_some() { - Some(2) - } else if storage.downcast_ref::().is_some() { - Some(4) - } else if let Some(integer) = storage.downcast_ref::() { - let width = u64::from(integer.width()); - (width > 0 && width % 8 == 0).then_some(width / 8) - } else { - None - } - } - - /// 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, - } - } -} - -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" - ); - } - let storage_ty = self.storage.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() - || logical_is_f32 - || plain_integer_is_supported; - if !plain_is_supported { - return verify_err!( - Location::Unknown, - "plain cute.tensor_view supports u8, u16, f16, and f32 carriers 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 => { - 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 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 cute.tensor_view must be read-only" - ); - } - if !matches!(self.role, CuteTensorRoleAttr::Mkl | CuteTensorRoleAttr::Nkl) { - return verify_err!( - Location::Unknown, - "UE8M0 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 cute.tensor_view must use BlockScaleKMajor" - ); - } - } - } - Ok(()) - } -} - -/// 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 { - pub tensor: TypeHandle, - 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.inner().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, - CuteTensorLayoutAttr::BlockScaleKMajor(_) - ) - ) { - return verify_err!( - Location::Unknown, - "cute.smem_tensor needs packed E2M1 values or UE8M0 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(()) - } -} - -/// 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 { - pub storage: TypeHandle, - 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) - } - - #[must_use] - pub fn half_elements(&self) -> Option { - self.plan.half_elements() - } - - #[must_use] - pub fn full_elements(&self) -> Option { - self.plan.full_elements() - } - - #[must_use] - pub fn half_bytes(&self) -> Option { - self.plan.half_bytes(2) - } - - #[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(()) - } -} - -/// 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 { - pub tensor: TypeHandle, - 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.inner().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(()) - } -} - -/// 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 { - pub values: TypeHandle, - pub scales: TypeHandle, - pub role: CuteTensorRoleAttr, - 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 scales.format != CuteTensorFormatAttr::UE8M0 - || scales.layout != CuteTensorLayoutAttr::BlockScaleKMajor(16) - { - return verify_err!( - Location::Unknown, - "cute.scaled_view scales must be a UE8M0 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 { - 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(()) - } -} - -/// 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 `WorkTile` struct reaches the -/// runtime ABI. -/// -/// ```text -/// current linear number + grid<16,16,1> -/// │ -/// ▼ -/// !cute.work_tile> -/// │ coordinates -/// ▼ -/// linear, m, n, batch -/// ``` -#[pliron_type(name = "cute.work_tile", format = "`<` $grid `>`")] -#[derive(Hash, PartialEq, Eq, Debug, Clone)] -pub struct CuteWorkTileType { - 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) - } -} - -/// 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 { - pub stages: u64, - pub consumer_warps: u32, - 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(()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::layout::{ComposedLayout, Layout, OffsetUnit, Swizzle}; - use pliron::builtin::types::Signedness; - - #[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()); - } - - #[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") - ); - } - - #[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)); - } - - #[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()); - } - - #[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/cuteir/dialect-cute/src/types/epilogue.rs b/cuda-oxide/cuteir/dialect-cute/src/types/epilogue.rs new file mode 100644 index 0000000000..6233e245ea --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/dialect-cute/src/types/gemv.rs b/cuda-oxide/cuteir/dialect-cute/src/types/gemv.rs new file mode 100644 index 0000000000..f9eece8be6 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/types/gemv.rs @@ -0,0 +1,203 @@ +/* + * 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 scales.format != CuteTensorFormatAttr::UE8M0 + || scales.layout != CuteTensorLayoutAttr::BlockScaleKMajor(16) + { + return verify_err!( + Location::Unknown, + "cute.scaled_view scales must be a UE8M0 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/cuteir/dialect-cute/src/types/mma.rs b/cuda-oxide/cuteir/dialect-cute/src/types/mma.rs new file mode 100644 index 0000000000..47265bb0de --- /dev/null +++ b/cuda-oxide/cuteir/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.inner().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, + CuteTensorLayoutAttr::BlockScaleKMajor(_) + ) + ) { + return verify_err!( + Location::Unknown, + "cute.smem_tensor needs packed E2M1 values or UE8M0 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/cuteir/dialect-cute/src/types/mod.rs b/cuda-oxide/cuteir/dialect-cute/src/types/mod.rs new file mode 100644 index 0000000000..e7feda4915 --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/dialect-cute/src/types/pipeline.rs b/cuda-oxide/cuteir/dialect-cute/src/types/pipeline.rs new file mode 100644 index 0000000000..f013857dce --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/dialect-cute/src/types/scheduler.rs b/cuda-oxide/cuteir/dialect-cute/src/types/scheduler.rs new file mode 100644 index 0000000000..19c2ae8006 --- /dev/null +++ b/cuda-oxide/cuteir/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 `WorkTile` struct reaches the +/// runtime ABI. +/// +/// ```text +/// current linear number + grid<16,16,1> +/// │ +/// ▼ +/// !cute.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/cuteir/dialect-cute/src/types/tensor.rs b/cuda-oxide/cuteir/dialect-cute/src/types/tensor.rs new file mode 100644 index 0000000000..4b105bd71d --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/types/tensor.rs @@ -0,0 +1,317 @@ +/* + * 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; +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, + ) + } + + /// Size of one v0 storage element in bytes. + #[must_use] + pub fn storage_bytes(&self, ctx: &Context) -> Option { + let storage = self.storage.deref(ctx); + if storage.downcast_ref::().is_some() { + Some(2) + } else if storage.downcast_ref::().is_some() { + Some(4) + } else if let Some(integer) = storage.downcast_ref::() { + let width = u64::from(integer.width()); + (width > 0 && width % 8 == 0).then_some(width / 8) + } else { + None + } + } + + /// 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, + } + } +} + +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" + ); + } + let storage_ty = self.storage.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() + || logical_is_f32 + || plain_integer_is_supported; + if !plain_is_supported { + return verify_err!( + Location::Unknown, + "plain cute.tensor_view supports u8, u16, f16, and f32 carriers 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 => { + 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 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 cute.tensor_view must be read-only" + ); + } + if !matches!(self.role, CuteTensorRoleAttr::Mkl | CuteTensorRoleAttr::Nkl) { + return verify_err!( + Location::Unknown, + "UE8M0 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 cute.tensor_view must use BlockScaleKMajor" + ); + } + } + } + Ok(()) + } +} diff --git a/cuda-oxide/cuteir/dialect-cute/src/types/tests.rs b/cuda-oxide/cuteir/dialect-cute/src/types/tests.rs new file mode 100644 index 0000000000..a542c77563 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/types/tests.rs @@ -0,0 +1,232 @@ +/* + * 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}; +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") + ); +} + +/// 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/cuteir/dialect-cute/src/types/tma.rs b/cuda-oxide/cuteir/dialect-cute/src/types/tma.rs new file mode 100644 index 0000000000..2add34a0a3 --- /dev/null +++ b/cuda-oxide/cuteir/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.inner().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/cuteir/dialect-cute/src/verify.rs b/cuda-oxide/cuteir/dialect-cute/src/verify.rs deleted file mode 100644 index 2312fd6fef..0000000000 --- a/cuda-oxide/cuteir/dialect-cute/src/verify.rs +++ /dev/null @@ -1,4548 +0,0 @@ -/* - * 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. - -use std::collections::{HashMap, HashSet}; - -use dialect_mir::attributes::MirCastKindAttr; -use dialect_mir::ops::control_flow::MirCondBranchOp; -use dialect_mir::ops::function::MirFuncOp; -use dialect_mir::ops::{ - MirAddOp, MirAllocaOp, MirCallOp, MirCastOp, MirConstantOp, MirConstructStructOp, - MirConstructTupleOp, MirExtractFieldOp, MirFieldAddrOp, MirGotoOp, MirInsertFieldOp, MirLoadOp, - MirNotOp, MirRefOp, MirStorageDeadOp, MirStorageLiveOp, MirStoreOp, MirUndefOp, -}; -use dialect_mir::types::{ - MirArrayType, MirDisjointSliceType, MirEnumType, MirPtrType, MirSliceType, MirStructType, - MirTupleType, MirUnionType, address_space, -}; -use pliron::attribute::attr_cast; -use pliron::basic_block::BasicBlock; -use pliron::builtin::attr_interfaces::TypedAttrInterface; -use pliron::builtin::op_interfaces::BranchOpInterface; -use pliron::builtin::type_interfaces::FunctionTypeInterface; -use pliron::builtin::types::IntegerType; -use pliron::common_traits::Verify; -use pliron::context::{Context, Ptr}; -use pliron::linked_list::ContainsLinkedList; -use pliron::op::{Op, OpId, op_cast}; -use pliron::operation::Operation; -use pliron::r#type::{TypeHandle, Typed, type_cast}; -use pliron::value::Value; - -use crate::attributes::{ - CuteEpilogueSyncPhaseAttr, CutePipelineRoleAttr, CutePipelineStateAttr, CuteTensorRoleAttr, - CuteTileGridAttr, CuteTiledMmaPlanAttr, -}; -use crate::epilogue_ops::{ - CuteEpilogueHalfOp, CuteEpilogueSmemOverlayOp, CuteEpilogueStoreFragmentOp, CuteEpilogueSyncOp, - CuteEpilogueWarpSliceOp, CuteTmaStore2dSemanticOp, CuteTmaStoreAcquireOp, CuteTmaStoreCommitOp, - CuteTmaStoreTailOp, -}; -use crate::gemm_tma_ops::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}; -use crate::gemv_ops::{ - CuteDotOp, CuteScaledViewKTileOp, CuteScaledViewLoadOp, CuteScaledViewMakeOp, - CuteScaledViewRowOp, CuteTensorMake2DOp, -}; -use crate::ops::{ - CuteAssumeDivOp, CuteCopyG2SOp, CuteCopyOp, CuteLdmatrixOp, CuteTmaLoad2dOp, CuteTmaStore2dOp, -}; -use crate::pipeline_ops::{ - CutePipelineConsumerReleaseOp, CutePipelineConsumerWaitOp, CutePipelineProducerAcquireOp, - CutePipelineProducerExpectTxOp, CutePipelineProducerTailOp, CutePipelineStateAdvanceOp, - CutePipelineStateNewOp, CutePipelineStateSlotOp, CuteTmaLoadPipelineInitOp, - CuteTmaLoadPipelineMakeOp, -}; -use crate::scheduler_ops::{ - CuteSchedulerAdvanceOp, CuteSchedulerCurrentOp, CuteSchedulerHasWorkOp, CuteSchedulerNew1dOp, - CuteWorkTileCoordinatesOp, -}; -use crate::smem_mma_ops::{ - CuteFragmentFillOp, CuteFragmentSliceKOp, CuteMmaLoadAOp, CuteMmaLoadScalesOp, - CuteMmaPartitionBOp, CuteSmemTensorOverlayOp, CuteTiledGemmOp, CuteTiledMmaSliceOp, -}; -use crate::tensor_ops::{ - CuteTensorBaseOp, CuteTensorIsFullOp, CuteTensorLoadIntoOp, CuteTensorMakeOp, - CuteTensorSliceOp, CuteTensorStoreElementAbsOp, CuteTensorStoreFromOp, - CuteTensorZippedDivideOp, -}; -use crate::types::{ - CuteEpilogueTileType, CuteFragmentType, CuteScaledViewType, CuteSmemTensorType, - CuteTensorViewType, CuteTmaLoadPipelineType, CuteTmaViewType, CuteWorkTileType, -}; - -/// 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 { - #[error("invalid CuTe semantic module: {0}")] - Invalid(String), - #[error("unsupported CuTe semantic operation: {0}")] - Unsupported(String), -} - -/// Result type shared by backend-neutral verification and backend adapters. -pub type VerifyResult = std::result::Result; - -type Result = VerifyResult; - -fn invalid(message: impl Into) -> VerifyError { - VerifyError::Invalid(message.into()) -} - -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); - } - } - } -} - -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(), - ] -} - -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(), - ] -} - -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(), - ] -} - -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(), - ] -} - -fn tma_semantic_ids() -> [OpId; 3] { - [ - CuteTmaGmemViewOp::get_opid_static(), - CuteTmaSmemViewOp::get_opid_static(), - CuteTmaCopy2dOp::get_opid_static(), - ] -} - -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(), - ] -} - -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(), - ] -} - -fn all_high_level_semantic_ids() -> Vec { - let mut ids = tensor_semantic_ids().to_vec(); - ids.push(crate::copy_ops::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::sm100_ops::semantic_ids()); - ids.extend(crate::sm100_attention_ops::semantic_ids()); - ids -} - -fn verify_local_op(ctx: &Context, operation: Ptr) -> Result { - let opid = Operation::get_opid(operation, ctx); - macro_rules! verify_as { - ($ty:ty) => { - <$ty>::wrap(operation).verify(ctx) - }; - } - - let result = if opid == crate::copy_ops::CuteStoreGlobalNoAllocateOp::get_opid_static() { - Some(verify_as!(crate::copy_ops::CuteStoreGlobalNoAllocateOp)) - } else if let Some(result) = crate::sm100_attention_ops::verify_local(ctx, operation) { - Some(result) - } else if let Some(result) = crate::sm100_ops::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(()) -} - -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 { - 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()) -} - -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(()) -} - -#[derive(Clone, Copy)] -struct TensorViewState { - tile_size: Option, - tile_index: Option, -} - -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}`" - ))) -} - -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(()) -} - -#[derive(Clone, Copy)] -struct GemvTensorState { - rows: Value, - k: Value, - role: CuteTensorRoleAttr, -} - -#[derive(Clone, Copy)] -struct GemvSelectionState { - values: GemvTensorState, - scales: GemvTensorState, - batch: Option, - row: Option, - tile_index: Option, -} - -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, - }) -} - -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}`" - ))) -} - -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(()) -} - -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(()) -} - -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(()) -} - -/// 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 -} - -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() -} - -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 -} - -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. -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)) -} - -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)) -} - -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -struct CarrierKey { - value: Value, - /// Empty means the value itself. Non-empty paths describe a field inside - /// an ordinary aggregate before a matching `mir.extract_field` exposes it. - path: Vec, -} - -#[derive(Clone)] -struct OrdinaryFacts { - facts: HashMap>, -} - -impl Default for OrdinaryFacts { - fn default() -> Self { - Self { - facts: HashMap::new(), - } - } -} - -impl OrdinaryFacts -where - F: Clone + Eq + std::hash::Hash, -{ - fn insert_path(&mut self, value: Value, path: Vec, fact: F) -> bool { - self.facts - .entry(CarrierKey { value, path }) - .or_default() - .insert(fact) - } - - fn insert(&mut self, value: Value, fact: F) -> bool { - self.insert_path(value, Vec::new(), fact) - } - - fn at_path(&self, value: Value, path: &[u32]) -> Option<&HashSet> { - self.facts.get(&CarrierKey { - value, - path: path.to_vec(), - }) - } - - fn at(&self, value: Value) -> Option<&HashSet> { - self.at_path(value, &[]) - } - - fn entries(&self, value: Value) -> Vec<(Vec, F)> { - self.facts - .iter() - .filter(|(key, _)| key.value == value) - .flat_map(|(key, facts)| facts.iter().cloned().map(|fact| (key.path.clone(), fact))) - .collect() - } - - 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 - } -} - -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) -} - -fn is_local_cell_pointer_alias_cast(ctx: &Context, operation: Ptr) -> bool { - is_provenance_preserving_pointer_cast(ctx, operation, false) -} - -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. -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. -fn local_cell_is_closed(ctx: &Context, root: Value) -> bool { - let mut worklist = vec![root]; - let mut visited = HashSet::new(); - while let Some(pointer) = worklist.pop() { - if !visited.insert(pointer) { - continue; - } - for r#use in pointer.uses(ctx) { - let user = r#use.user_op(); - let opid = Operation::get_opid(user, ctx); - let operation = user.deref(ctx); - if r#use.find_index(ctx) != 0 { - return false; - } - let forwards_pointer = operation.get_num_results() == 1 - && ((opid == MirCastOp::get_opid_static() - && is_local_cell_pointer_alias_cast(ctx, user)) - || (opid == MirFieldAddrOp::get_opid_static() - && MirFieldAddrOp::new(user) - .get_attr_field_index(ctx) - .is_some())); - if forwards_pointer { - worklist.push(operation.get_result(0)); - } else if opid == MirLoadOp::get_opid_static() { - if MirLoadOp::new(user).is_volatile(ctx) { - return false; - } - } else if opid == MirStoreOp::get_opid_static() { - if MirStoreOp::new(user).is_volatile(ctx) { - return false; - } - } else { - return false; - } - } - } - true -} - -/// Propagate semantic facts through the exact ordinary SSA adapters that can -/// remain after common CuTe preparation. This is deliberately a monotone -/// fixed-point transfer: loop block arguments become known once their entry -/// fact reaches them, then semantic advance operations can feed the same fact -/// around the backedge on a later iteration. -fn propagate_ordinary_plumbing( - ctx: &Context, - operations: &[Ptr], - facts: &mut OrdinaryFacts, -) -> bool -where - F: Clone + Eq + std::hash::Hash, -{ - let mut changed = false; - let mut blocks = HashSet::new(); - for operation in operations { - let operation_ref = operation.deref(ctx); - if let Some(block) = operation_ref.get_parent_block() { - blocks.insert(block); - } - let opid = Operation::get_opid(*operation, ctx); - if opid == MirCastOp::get_opid_static() - && is_semantic_pointer_carrier_cast(ctx, *operation) - && operation_ref.get_num_operands() == 1 - && operation_ref.get_num_results() == 1 - { - changed |= facts.copy_entries( - operation_ref.get_operand(0), - operation_ref.get_result(0), - &[], - &[], - ); - } else if (opid == MirConstructStructOp::get_opid_static() - || opid == MirConstructTupleOp::get_opid_static()) - && operation_ref.get_num_results() == 1 - { - let result = operation_ref.get_result(0); - for (index, operand) in operation_ref.operands().enumerate() { - changed |= facts.copy_entries(operand, result, &[], &[index as u32]); - } - } else if opid == MirInsertFieldOp::get_opid_static() - && operation_ref.get_num_operands() == 2 - && operation_ref.get_num_results() == 1 - { - let Some(index) = MirInsertFieldOp::new(*operation).get_attr_insert_index(ctx) else { - continue; - }; - let base = operation_ref.get_operand(0); - let inserted = operation_ref.get_operand(1); - let result = operation_ref.get_result(0); - for (path, fact) in facts.entries(base) { - if path.first().copied() != Some(index.0) { - changed |= facts.insert_path(result, path, fact); - } - } - changed |= facts.copy_entries(inserted, result, &[], &[index.0]); - } else if opid == MirExtractFieldOp::get_opid_static() - && operation_ref.get_num_operands() == 1 - && operation_ref.get_num_results() == 1 - { - let Some(index) = MirExtractFieldOp::new(*operation).get_attr_index(ctx) else { - continue; - }; - changed |= facts.copy_entries( - operation_ref.get_operand(0), - operation_ref.get_result(0), - &[index.0], - &[], - ); - } else if opid == MirStoreOp::get_opid_static() && operation_ref.get_num_operands() == 2 { - let store = MirStoreOp::new(*operation); - if !store.is_volatile(ctx) - && let Some((root, path)) = local_cell(ctx, store.address_opd(ctx)) - && local_cell_is_closed(ctx, root) - { - changed |= facts.copy_entries(store.value_opd(ctx), root, &[], &path); - } - } else if opid == MirLoadOp::get_opid_static() - && operation_ref.get_num_operands() == 1 - && operation_ref.get_num_results() == 1 - { - let load = MirLoadOp::new(*operation); - if !load.is_volatile(ctx) - && let Some((root, path)) = local_cell(ctx, load.address_opd(ctx)) - && local_cell_is_closed(ctx, root) - { - changed |= facts.copy_entries(root, operation_ref.get_result(0), &path, &[]); - } - } - } - - for block in blocks { - let block_ref = block.deref(ctx); - for argument_index in 0..block_ref.get_num_arguments() { - let argument = block_ref.get_argument(argument_index); - for edge in block.uses(ctx) { - let terminator = edge.user_op(); - let Some(operands) = - branch_successor_operands(ctx, terminator, edge.find_index(ctx)) - else { - continue; - }; - if let Some(incoming) = operands.get(argument_index).copied() { - changed |= facts.copy_entries(incoming, argument, &[], &[]); - } - } - } - } - changed -} - -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(); - } - } -} - -fn seed_unknown_integer_carriers( - ctx: &Context, - operations: &[Ptr], - facts: &mut OrdinaryFacts, - width: u32, - unknown: impl Fn(Value) -> F, -) where - F: Clone + Eq + std::hash::Hash, -{ - let mut values = Vec::new(); - let mut blocks = HashSet::new(); - for operation in operations { - let operation_ref = operation.deref(ctx); - values.extend(operation_ref.results()); - if let Some(block) = operation_ref.get_parent_block() - && blocks.insert(block) - { - let block_ref = block.deref(ctx); - values.extend( - (0..block_ref.get_num_arguments()).map(|index| block_ref.get_argument(index)), - ); - } - } - - for value in values { - let mut paths = Vec::new(); - integer_leaf_paths(ctx, value.get_type(ctx), width, &mut Vec::new(), &mut paths); - for path in paths { - if facts.at_path(value, &path).is_none() { - facts.insert_path(value, path, unknown(value)); - } - } - } -} - -fn poison_missing_forwarding_sources( - ctx: &Context, - operations: &[Ptr], - facts: &mut OrdinaryFacts, - unknown: impl Fn(Value) -> F + Copy, -) -> bool -where - F: Clone + Eq + std::hash::Hash, -{ - let mut changed = false; - let mut blocks = HashSet::new(); - for operation in operations { - if let Some(block) = operation.deref(ctx).get_parent_block() { - blocks.insert(block); - } - if Operation::get_opid(*operation, ctx) != MirStoreOp::get_opid_static() { - continue; - } - let store = MirStoreOp::new(*operation); - let Some((root, cell_path)) = local_cell(ctx, store.address_opd(ctx)) else { - continue; - }; - if !local_cell_is_closed(ctx, root) { - continue; - } - for (path, _) in facts.entries(root) { - let Some(source_path) = path.strip_prefix(cell_path.as_slice()) else { - continue; - }; - if facts.at_path(store.value_opd(ctx), source_path).is_none() { - changed |= facts.insert_path( - store.value_opd(ctx), - source_path.to_vec(), - unknown(store.value_opd(ctx)), - ); - } - } - } - - for block in blocks { - let block_ref = block.deref(ctx); - for argument_index in 0..block_ref.get_num_arguments() { - let argument = block_ref.get_argument(argument_index); - let argument_entries = facts.entries(argument); - if argument_entries.is_empty() { - continue; - } - for edge in block.uses(ctx) { - let terminator = edge.user_op(); - let Some(operands) = - branch_successor_operands(ctx, terminator, edge.find_index(ctx)) - else { - continue; - }; - let Some(incoming) = operands.get(argument_index).copied() else { - continue; - }; - for (path, _) in &argument_entries { - if facts.at_path(incoming, path).is_none() { - changed |= facts.insert_path(incoming, path.clone(), unknown(incoming)); - } - } - } - } - } - changed -} - -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"))) -} - -#[derive(Clone, Copy, PartialEq, Eq, Hash)] -enum SchedulerFact { - Current(Ptr), - Stride(Ptr), - Unknown(Value), -} - -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" - ))), - } -} - -fn build_scheduler_facts( - ctx: &Context, - operations: &[Ptr], -) -> OrdinaryFacts { - let mut facts = OrdinaryFacts::default(); - for operation in operations { - if Operation::get_opid(*operation, ctx) == CuteSchedulerNew1dOp::get_opid_static() { - let scheduler = CuteSchedulerNew1dOp::wrap(*operation); - facts.insert(scheduler.current(ctx), SchedulerFact::Current(*operation)); - facts.insert(scheduler.stride(ctx), SchedulerFact::Stride(*operation)); - } - } - - loop { - let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); - for operation in operations { - if Operation::get_opid(*operation, ctx) != CuteSchedulerAdvanceOp::get_opid_static() { - continue; - } - let advance = CuteSchedulerAdvanceOp::wrap(*operation); - let Some(current_facts) = facts.at(advance.current(ctx)).cloned() else { - continue; - }; - let Some(stride_facts) = facts.at(advance.stride(ctx)).cloned() else { - continue; - }; - for current in ¤t_facts { - let SchedulerFact::Current(root) = current else { - continue; - }; - if stride_facts.contains(&SchedulerFact::Stride(*root)) { - changed |= facts.insert(advance.next(ctx), SchedulerFact::Current(*root)); - } - } - } - if !changed { - break; - } - } - - // A factless scalar is an ordinary, unrelated value. Seeding it only - // after the semantic fixed point lets a valid loop recurrence establish - // itself first, while still poisoning a merge with an arbitrary edge. - seed_unknown_integer_carriers(ctx, operations, &mut facts, 64, SchedulerFact::Unknown); - loop { - let poisoned = - poison_missing_forwarding_sources(ctx, operations, &mut facts, SchedulerFact::Unknown); - let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); - if !poisoned && !propagated { - break; - } - } - facts -} - -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), - )); - } -} - -fn verify_scheduler_story(ctx: &Context, all_ops: &[Ptr]) -> Result { - let ids = scheduler_semantic_ids(); - let facts = build_scheduler_facts(ctx, all_ops); - let mut roots = HashSet::new(); - let mut has_work_roots = HashSet::new(); - let mut current_roots = HashSet::new(); - let mut advance_roots = HashSet::new(); - let mut coordinate_roots = HashSet::new(); - let mut has_work_by_root: HashMap, Vec>> = HashMap::new(); - let mut current_by_root: HashMap, Vec>> = HashMap::new(); - let mut advance_by_root: HashMap, Vec>> = HashMap::new(); - let mut coordinates_by_root: HashMap, Vec>> = HashMap::new(); - let mut currents = Vec::new(); - let mut coordinates = Vec::new(); - for operation in all_ops { - let opid = Operation::get_opid(*operation, ctx); - if !ids.contains(&opid) { - continue; - } - if operation.deref(ctx).get_parent_block().is_none() { - return Err(invalid(format!( - "scheduler operation `{opid}` is not attached to a block" - ))); - } - let grid = scheduler_grid(ctx, *operation, &opid)?; - if opid == CuteSchedulerNew1dOp::get_opid_static() { - roots.insert(*operation); - } else if opid == CuteSchedulerHasWorkOp::get_opid_static() { - let semantic = CuteSchedulerHasWorkOp::wrap(*operation); - let root = scheduler_root( - &facts, - semantic.current(ctx), - true, - "cute.scheduler_has_work current", - )?; - if CuteSchedulerNew1dOp::wrap(root).tile_grid(ctx) != Some(grid) { - return Err(invalid( - "cute.scheduler_has_work grid differs from its scheduler_new_1d root", - )); - } - if !semantic.has_work(ctx).is_used(ctx) { - return Err(invalid( - "cute.scheduler_has_work result is unused; it must guard the scheduler body", - )); - } - has_work_roots.insert(root); - has_work_by_root.entry(root).or_default().push(*operation); - } else if opid == CuteSchedulerCurrentOp::get_opid_static() { - let semantic = CuteSchedulerCurrentOp::wrap(*operation); - let root = scheduler_root( - &facts, - semantic.current(ctx), - true, - "cute.scheduler_current current", - )?; - if CuteSchedulerNew1dOp::wrap(root).tile_grid(ctx) != Some(grid) { - return Err(invalid( - "cute.scheduler_current grid differs from its scheduler_new_1d root", - )); - } - current_roots.insert(root); - current_by_root.entry(root).or_default().push(*operation); - currents.push(*operation); - } else if opid == CuteWorkTileCoordinatesOp::get_opid_static() { - coordinates.push(*operation); - } else if opid == CuteSchedulerAdvanceOp::get_opid_static() { - let semantic = CuteSchedulerAdvanceOp::wrap(*operation); - let current_root = scheduler_root( - &facts, - semantic.current(ctx), - true, - "cute.scheduler_advance current", - )?; - let stride_root = scheduler_root( - &facts, - semantic.stride(ctx), - false, - "cute.scheduler_advance stride", - )?; - if current_root != stride_root - || CuteSchedulerNew1dOp::wrap(current_root).tile_grid(ctx) != Some(grid) - { - return Err(invalid( - "cute.scheduler_advance must use current, stride, and grid from one scheduler_new_1d", - )); - } - advance_roots.insert(current_root); - advance_by_root - .entry(current_root) - .or_default() - .push(*operation); - } - } - - for current in ¤ts { - let tile = CuteSchedulerCurrentOp::wrap(*current).work_tile(ctx); - let uses = tile.uses(ctx); - if uses.is_empty() - || uses.iter().any(|r#use| { - Operation::get_opid(r#use.user_op(), ctx) - != CuteWorkTileCoordinatesOp::get_opid_static() - || r#use.find_index(ctx) != 0 - }) - { - return Err(invalid( - "cute.scheduler_current work tile must be used only by at least one cute.work_tile_coordinates", - )); - } - } - for coordinates in coordinates { - let tile = CuteWorkTileCoordinatesOp::wrap(coordinates).work_tile(ctx); - let Some(producer) = tile.defining_op() else { - return Err(invalid( - "cute.work_tile_coordinates reached a block argument; a direct scheduler_current result is required", - )); - }; - if Operation::get_opid(producer, ctx) != CuteSchedulerCurrentOp::get_opid_static() - || !currents.contains(&producer) - { - return Err(invalid( - "cute.work_tile_coordinates must use a scheduler_current attached to this module", - )); - } - let root = scheduler_root( - &facts, - CuteSchedulerCurrentOp::wrap(producer).current(ctx), - true, - "cute.work_tile_coordinates scheduler current", - )?; - coordinate_roots.insert(root); - coordinates_by_root - .entry(root) - .or_default() - .push(coordinates); - } - - for root in roots { - if !has_work_roots.contains(&root) - || !current_roots.contains(&root) - || !coordinate_roots.contains(&root) - || !advance_roots.contains(&root) - { - return Err(invalid( - "each cute.scheduler_new_1d needs a complete has_work -> current -> coordinates -> advance story", - )); - } - let has_work = &has_work_by_root[&root]; - let current = ¤t_by_root[&root]; - let coordinates = &coordinates_by_root[&root]; - let advances = &advance_by_root[&root]; - if has_work.len() != 1 - || current.len() != 1 - || coordinates.len() != 1 - || advances.len() != 1 - { - return Err(invalid( - "each scheduler root needs exactly one static has_work/current/coordinates/advance protocol", - )); - } - let Some((guard, work_successor, exit_successor)) = - scheduler_guard(ctx, CuteSchedulerHasWorkOp::wrap(has_work[0]).has_work(ctx)) - else { - return Err(invalid( - "cute.scheduler_has_work must exclusively guard one conditional branch (optionally through mir.not)", - )); - }; - let current_block = current[0] - .deref(ctx) - .get_parent_block() - .expect("attached scheduler current"); - if !operation_reaches(ctx, root, has_work[0]) - || !block_reaches(ctx, work_successor, current_block) - || block_reaches(ctx, exit_successor, current_block) - || !operation_reaches(ctx, current[0], coordinates[0]) - || !operation_reaches(ctx, coordinates[0], advances[0]) - || !operation_reaches(ctx, advances[0], guard) - { - return Err(invalid( - "scheduler CFG must be new -> guarded work current -> coordinates -> advance -> next guard, with the exit edge outside the work body", - )); - } - } - Ok(()) -} - -#[derive(Clone, Copy)] -struct LoadPipelineFacts { - make: Ptr, - stages: u64, - transaction_bytes: u32, -} - -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, - }) -} - -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 - } -} - -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) -} - -#[derive(Clone, Copy, PartialEq, Eq, Hash)] -enum PipelineFact { - Slot(Ptr), - Phase(Ptr), - Unknown(Value), -} - -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" - ))), - } -} - -fn build_pipeline_facts( - ctx: &Context, - operations: &[Ptr], -) -> OrdinaryFacts { - let mut facts = OrdinaryFacts::default(); - for operation in operations { - if Operation::get_opid(*operation, ctx) == CutePipelineStateNewOp::get_opid_static() { - let state = CutePipelineStateNewOp::wrap(*operation); - facts.insert(state.slot(ctx), PipelineFact::Slot(*operation)); - facts.insert(state.phase(ctx), PipelineFact::Phase(*operation)); - } - } - - loop { - let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); - for operation in operations { - if Operation::get_opid(*operation, ctx) != CutePipelineStateAdvanceOp::get_opid_static() - { - continue; - } - let advance = CutePipelineStateAdvanceOp::wrap(*operation); - let Some(slot_facts) = facts.at(advance.slot(ctx)).cloned() else { - continue; - }; - let Some(phase_facts) = facts.at(advance.phase(ctx)).cloned() else { - continue; - }; - for slot in &slot_facts { - let PipelineFact::Slot(root) = slot else { - continue; - }; - if phase_facts.contains(&PipelineFact::Phase(*root)) - && CutePipelineStateNewOp::wrap(*root).state(ctx) == advance.state(ctx) - { - changed |= facts.insert(advance.next_slot(ctx), PipelineFact::Slot(*root)); - changed |= facts.insert(advance.next_phase(ctx), PipelineFact::Phase(*root)); - } - } - } - if !changed { - break; - } - } - - seed_unknown_integer_carriers(ctx, operations, &mut facts, 32, PipelineFact::Unknown); - loop { - let poisoned = - poison_missing_forwarding_sources(ctx, operations, &mut facts, PipelineFact::Unknown); - let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); - if !poisoned && !propagated { - break; - } - } - facts -} - -#[derive(Default)] -struct PipelineStateEvents { - slots: Vec>, - advances: Vec>, - acquires: Vec>, - expects: Vec>, - waits: Vec>, - releases: Vec>, - tails: Vec>, -} - -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", - )) - } -} - -fn verify_pipeline_story(ctx: &Context, all_ops: &[Ptr]) -> Result { - let ids = pipeline_semantic_ids(); - let facts = build_pipeline_facts(ctx, all_ops); - let makes: Vec<_> = all_ops - .iter() - .copied() - .filter(|operation| { - Operation::get_opid(*operation, ctx) == CuteTmaLoadPipelineMakeOp::get_opid_static() - }) - .collect(); - let mut init_by_make: HashMap, Vec>> = HashMap::new(); - let mut events: HashMap, PipelineStateEvents> = HashMap::new(); - let mut binding: HashMap, Ptr> = HashMap::new(); - - for make in &makes { - if make.deref(ctx).get_parent_block().is_none() { - return Err(invalid("TMA load-pipeline make is not attached to a block")); - } - let handle = CuteTmaLoadPipelineMakeOp::wrap(*make).pipeline(ctx); - let uses = handle.uses(ctx); - if uses.is_empty() { - return Err(invalid( - "cute.tma_load_pipeline_make has no lifecycle consumer", - )); - } - for r#use in uses { - let user = r#use.user_op(); - let user_id = Operation::get_opid(user, ctx); - if !ids.contains(&user_id) - || r#use.find_index(ctx) != 0 - || pipeline_operand(ctx, user, &user_id) != Some(handle) - { - return Err(invalid(format!( - "TMA load-pipeline handle has an invalid use by `{user_id}`" - ))); - } - } - } - - for operation in all_ops { - let opid = Operation::get_opid(*operation, ctx); - if !ids.contains(&opid) { - continue; - } - if operation.deref(ctx).get_parent_block().is_none() { - return Err(invalid(format!( - "pipeline operation `{opid}` is not attached to a block" - ))); - } - if opid == CutePipelineStateNewOp::get_opid_static() { - events.entry(*operation).or_default(); - continue; - } - - let state_root = if opid == CutePipelineStateSlotOp::get_opid_static() { - let semantic = CutePipelineStateSlotOp::wrap(*operation); - let root = pipeline_state_root( - &facts, - semantic.slot(ctx), - true, - "cute.pipeline_state_slot slot", - )?; - if semantic.state(ctx) != CutePipelineStateNewOp::wrap(root).state(ctx) { - return Err(invalid( - "cute.pipeline_state_slot facts differ from its pipeline_state_new root", - )); - } - events.entry(root).or_default().slots.push(*operation); - Some(root) - } else if opid == CutePipelineStateAdvanceOp::get_opid_static() { - let semantic = CutePipelineStateAdvanceOp::wrap(*operation); - let slot_root = pipeline_state_root( - &facts, - semantic.slot(ctx), - true, - "cute.pipeline_state_advance slot", - )?; - let phase_root = pipeline_state_root( - &facts, - semantic.phase(ctx), - false, - "cute.pipeline_state_advance phase", - )?; - if slot_root != phase_root - || semantic.state(ctx) != CutePipelineStateNewOp::wrap(slot_root).state(ctx) - { - return Err(invalid( - "cute.pipeline_state_advance must consume slot and phase from one matching pipeline_state_new", - )); - } - events - .entry(slot_root) - .or_default() - .advances - .push(*operation); - Some(slot_root) - } else if opid == CuteTmaLoadPipelineMakeOp::get_opid_static() { - None - } else if opid == CuteTmaLoadPipelineInitOp::get_opid_static() { - let pipeline = resolve_load_pipeline( - ctx, - CuteTmaLoadPipelineInitOp::wrap(*operation).pipeline(ctx), - )?; - if !makes.contains(&pipeline.make) { - return Err(invalid( - "pipeline init refers to a make outside its semantic scope", - )); - } - init_by_make - .entry(pipeline.make) - .or_default() - .push(*operation); - None - } else { - let state = pipeline_state_for_op(ctx, *operation, &opid) - .ok_or_else(|| invalid(format!("pipeline operation `{opid}` has no state")))?; - let operation_ref = operation.deref(ctx); - let slot_root = pipeline_state_root( - &facts, - operation_ref.get_operand(1), - true, - &format!("`{opid}` slot"), - )?; - let needs_phase = opid == CutePipelineProducerAcquireOp::get_opid_static() - || opid == CutePipelineConsumerWaitOp::get_opid_static() - || opid == CutePipelineProducerTailOp::get_opid_static(); - if needs_phase { - let phase_root = pipeline_state_root( - &facts, - operation_ref.get_operand(2), - false, - &format!("`{opid}` phase"), - )?; - if phase_root != slot_root { - return Err(invalid(format!( - "`{opid}` mixes slot and phase from different pipeline states" - ))); - } - } - if CutePipelineStateNewOp::wrap(slot_root).state(ctx) != Some(state) { - return Err(invalid(format!( - "`{opid}` state facts differ from its pipeline_state_new root" - ))); - } - - let handle = pipeline_operand(ctx, *operation, &opid) - .ok_or_else(|| invalid(format!("`{opid}` has no load-pipeline handle")))?; - let pipeline = resolve_load_pipeline(ctx, handle)?; - if !makes.contains(&pipeline.make) { - return Err(invalid( - "pipeline lifecycle operation refers to a make outside its semantic scope", - )); - } - if state.stages != pipeline.stages { - return Err(invalid(format!( - "pipeline state on `{opid}` has {} stages but its handle has {}", - state.stages, pipeline.stages - ))); - } - if let Some(previous) = binding.insert(slot_root, pipeline.make) - && previous != pipeline.make - { - return Err(invalid( - "one ordinary pipeline state is used with multiple load-pipeline handles", - )); - } - - let root_events = events.entry(slot_root).or_default(); - if opid == CutePipelineProducerAcquireOp::get_opid_static() { - root_events.acquires.push(*operation); - } else if opid == CutePipelineProducerExpectTxOp::get_opid_static() { - root_events.expects.push(*operation); - let copied = expected_tma_transaction_bytes(ctx, *operation)?; - if copied != u64::from(pipeline.transaction_bytes) { - return Err(invalid(format!( - "cute.pipeline_producer_expect_tx promises {} bytes, but its direct TMA copies move {copied} bytes", - pipeline.transaction_bytes - ))); - } - } else if opid == CutePipelineConsumerWaitOp::get_opid_static() { - root_events.waits.push(*operation); - } else if opid == CutePipelineConsumerReleaseOp::get_opid_static() { - root_events.releases.push(*operation); - } else if opid == CutePipelineProducerTailOp::get_opid_static() { - root_events.tails.push(*operation); - } - Some(slot_root) - }; - - if let Some(root) = state_root { - let root_state = CutePipelineStateNewOp::wrap(root) - .state(ctx) - .expect("locally verified pipeline state root"); - if let Some(state) = pipeline_state_for_op(ctx, *operation, &opid) - && state != root_state - { - return Err(invalid(format!( - "`{opid}` does not carry its pipeline-state root's role and stage count" - ))); - } - } - } - - let mut roles_by_make: HashMap, (usize, usize)> = HashMap::new(); - for (root, root_events) in &events { - let state = CutePipelineStateNewOp::wrap(*root) - .state(ctx) - .expect("locally verified pipeline state root"); - let Some(make) = binding.get(root).copied() else { - return Err(invalid( - "cute.pipeline_state_new is not connected to any load-pipeline lifecycle", - )); - }; - if root_events.slots.is_empty() || root_events.advances.is_empty() { - return Err(invalid( - "each pipeline state needs state_slot and state_advance operations", - )); - } - let roles = roles_by_make.entry(make).or_default(); - match state.role { - CutePipelineRoleAttr::Producer => { - roles.0 += 1; - if root_events.acquires.len() != 1 - || root_events.expects.len() != 1 - || root_events.advances.len() != 1 - || root_events.slots.len() != 1 - || root_events.tails.len() != 1 - || !root_events.waits.is_empty() - || !root_events.releases.is_empty() - { - return Err(invalid( - "producer pipeline state needs acquire, expect_tx, advance, and exactly one tail, with no consumer events", - )); - } - let acquire = root_events.acquires[0]; - let expect = root_events.expects[0]; - let advance = root_events.advances[0]; - let tail = root_events.tails[0]; - require_producer_pipeline_order(ctx, *root, acquire, expect, advance, tail)?; - let barrier = CutePipelineProducerExpectTxOp::wrap(expect).completion_barrier(ctx); - for r#use in barrier.uses(ctx) { - let copy = r#use.user_op(); - if !operation_dominates(ctx, expect, copy) - || !operation_dominates(ctx, copy, advance) - { - return Err(invalid( - "every TMA copy must be CFG-ordered after expect_tx and before producer state advance", - )); - } - } - } - CutePipelineRoleAttr::Consumer => { - roles.1 += 1; - if root_events.waits.len() != 1 - || root_events.releases.len() != 1 - || root_events.advances.len() != 1 - || root_events.slots.len() != 1 - || !root_events.acquires.is_empty() - || !root_events.expects.is_empty() - || !root_events.tails.is_empty() - { - return Err(invalid( - "consumer pipeline state needs wait, release, and advance, with no producer events", - )); - } - let wait = root_events.waits[0]; - let release = root_events.releases[0]; - let advance = root_events.advances[0]; - if !operation_dominates(ctx, *root, wait) - || !operation_dominates(ctx, wait, release) - || !operation_dominates(ctx, release, advance) - || !operation_reaches(ctx, advance, wait) - { - return Err(invalid( - "consumer pipeline CFG must order state_new -> wait -> release -> advance -> next wait", - )); - } - } - } - } - for make in makes { - if init_by_make - .get(&make) - .is_none_or(|initializers| initializers.len() != 1) - { - return Err(invalid( - "each tma_load_pipeline_make needs exactly one initialization", - )); - } - if roles_by_make.get(&make).copied() != Some((1, 1)) { - return Err(invalid( - "each tma_load_pipeline_make needs exactly one complete producer state and one complete consumer state", - )); - } - let initializer = init_by_make[&make][0]; - for (root, bound_make) in &binding { - if *bound_make == make && !operation_dominates(ctx, initializer, *root) { - return Err(invalid( - "load-pipeline initialization must dominate both producer and consumer state creation", - )); - } - } - } - Ok(()) -} - -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) -} - -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(()) -} - -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(()) -} - -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 - } -} - -#[derive(Clone, Copy, PartialEq, Eq, Hash)] -enum MmaFact { - Lane(Ptr), - SmemBase(Ptr), - SmemCapacity(Ptr), - StageScales(Ptr), - SlicedScales(Ptr), - AFragment(Ptr), - BBase(Ptr), - BCapacity(Ptr), - BWarpN(Ptr), - BKHalf(Ptr), - Accumulator(Ptr), - Scalar(Value), - Unknown(Value), -} - -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")) -} - -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" - ))), - } -} - -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) -} - -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" - ))), - } -} - -fn build_mma_facts(ctx: &Context, operations: &[Ptr]) -> OrdinaryFacts { - let mut facts = OrdinaryFacts::default(); - for operation in operations { - let opid = Operation::get_opid(*operation, ctx); - let operation_ref = operation.deref(ctx); - if opid == CuteSmemTensorOverlayOp::get_opid_static() { - let overlay = CuteSmemTensorOverlayOp::wrap(*operation); - facts.insert(overlay.base(ctx), MmaFact::SmemBase(*operation)); - facts.insert(overlay.capacity(ctx), MmaFact::SmemCapacity(*operation)); - } else if opid == CuteTiledMmaSliceOp::get_opid_static() { - facts.insert( - CuteTiledMmaSliceOp::wrap(*operation).sliced_lane(ctx), - MmaFact::Lane(*operation), - ); - } else if opid == CuteFragmentFillOp::get_opid_static() { - facts.insert( - CuteFragmentFillOp::wrap(*operation).fragment(ctx), - MmaFact::Accumulator(*operation), - ); - } else if opid == CuteMmaLoadScalesOp::get_opid_static() { - facts.insert( - CuteMmaLoadScalesOp::wrap(*operation).scales(ctx), - MmaFact::StageScales(*operation), - ); - } else if opid == CuteFragmentSliceKOp::get_opid_static() { - facts.insert( - CuteFragmentSliceKOp::wrap(*operation).scales(ctx), - MmaFact::SlicedScales(*operation), - ); - } else if opid == CuteMmaLoadAOp::get_opid_static() { - facts.insert( - CuteMmaLoadAOp::wrap(*operation).fragment(ctx), - MmaFact::AFragment(*operation), - ); - } else if opid == CuteMmaPartitionBOp::get_opid_static() { - facts.insert(operation_ref.get_result(0), MmaFact::BBase(*operation)); - facts.insert(operation_ref.get_result(1), MmaFact::BCapacity(*operation)); - facts.insert(operation_ref.get_result(2), MmaFact::BWarpN(*operation)); - facts.insert(operation_ref.get_result(3), MmaFact::BKHalf(*operation)); - } - } - - loop { - let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); - for operation in operations { - if Operation::get_opid(*operation, ctx) != CuteTiledGemmOp::get_opid_static() { - continue; - } - let operation_ref = operation.deref(ctx); - let Some(accumulator_facts) = facts.at(operation_ref.get_operand(7)).cloned() else { - continue; - }; - for accumulator in accumulator_facts { - if let MmaFact::Accumulator(root) = accumulator { - changed |= facts.insert( - CuteTiledGemmOp::wrap(*operation).accumulator(ctx), - MmaFact::Accumulator(root), - ); - } - } - } - if !changed { - break; - } - } - - seed_unknown_integer_carriers(ctx, operations, &mut facts, 32, MmaFact::Scalar); - seed_unknown_integer_carriers(ctx, operations, &mut facts, 64, MmaFact::Scalar); - loop { - let poisoned = - poison_missing_forwarding_sources(ctx, operations, &mut facts, MmaFact::Unknown); - let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); - if !poisoned && !propagated { - break; - } - } - facts -} - -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)) -} - -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, - } -} - -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, - ) - }) - }) - }) -} - -fn mma_ordinary_use_forwards_fact( - ctx: &Context, - facts: &OrdinaryFacts, - key: &CarrierKey, - fact: MmaFact, - user: Ptr, - operand_index: usize, -) -> bool { - let opid = Operation::get_opid(user, ctx); - let operation = user.deref(ctx); - if opid == MirCastOp::get_opid_static() - && operand_index == 0 - && is_semantic_pointer_carrier_cast(ctx, user) - && operation.get_num_results() == 1 - { - return mma_fact_at_path(facts, operation.get_result(0), &key.path, fact); - } - if (opid == MirConstructStructOp::get_opid_static() - || opid == MirConstructTupleOp::get_opid_static()) - && operation.get_num_results() == 1 - { - let mut destination_path = vec![operand_index as u32]; - destination_path.extend_from_slice(&key.path); - return mma_fact_at_path(facts, operation.get_result(0), &destination_path, fact); - } - if opid == MirInsertFieldOp::get_opid_static() - && operation.get_num_operands() == 2 - && operation.get_num_results() == 1 - { - let Some(index) = MirInsertFieldOp::new(user).get_attr_insert_index(ctx) else { - return false; - }; - if operand_index == 0 { - if key.path.first().copied() == Some(index.0) { - return true; - } - return mma_fact_at_path(facts, operation.get_result(0), &key.path, fact); - } - if operand_index == 1 { - let mut destination_path = vec![index.0]; - destination_path.extend_from_slice(&key.path); - return mma_fact_at_path(facts, operation.get_result(0), &destination_path, fact); - } - return false; - } - if opid == MirExtractFieldOp::get_opid_static() - && operand_index == 0 - && operation.get_num_results() == 1 - { - let Some(index) = MirExtractFieldOp::new(user).get_attr_index(ctx) else { - return false; - }; - if key.path.first().copied() != Some(index.0) { - return true; - } - return mma_fact_at_path(facts, operation.get_result(0), &key.path[1..], fact); - } - if opid == MirStoreOp::get_opid_static() && operation.get_num_operands() == 2 { - let store = MirStoreOp::new(user); - if operand_index == 0 { - return local_cell(ctx, store.address_opd(ctx)) - .is_some_and(|(root, _)| local_cell_is_closed(ctx, root)); - } - if operand_index == 1 - && let Some((root, mut destination_path)) = local_cell(ctx, store.address_opd(ctx)) - && local_cell_is_closed(ctx, root) - { - destination_path.extend_from_slice(&key.path); - return mma_fact_at_path(facts, root, &destination_path, fact); - } - return false; - } - if opid == MirLoadOp::get_opid_static() - && operand_index == 0 - && operation.get_num_results() == 1 - { - let load = MirLoadOp::new(user); - let Some((root, cell_path)) = local_cell(ctx, load.address_opd(ctx)) else { - return false; - }; - if !local_cell_is_closed(ctx, root) { - return false; - } - let Some(result_path) = key.path.strip_prefix(cell_path.as_slice()) else { - return true; - }; - return mma_fact_at_path(facts, operation.get_result(0), result_path, fact); - } - if opid == MirFieldAddrOp::get_opid_static() - && operand_index == 0 - && operation.get_num_results() == 1 - && MirFieldAddrOp::new(user) - .get_attr_field_index(ctx) - .is_some() - { - return local_cell(ctx, operation.get_result(0)) - .is_some_and(|(root, _)| local_cell_is_closed(ctx, root)); - } - if opid == MirCondBranchOp::get_opid_static() && operand_index == 0 { - return false; - } - mma_branch_forwards_fact(ctx, facts, user, key.value, &key.path, fact) -} - -fn verify_mma_use_closure( - ctx: &Context, - operations: &[Ptr], - facts: &OrdinaryFacts, -) -> Result { - for (key, found) in &facts.facts { - for fact in found.iter().copied() { - if matches!(fact, MmaFact::Scalar(_) | MmaFact::Unknown(_)) { - continue; - } - for r#use in key.value.uses(ctx) { - let user = r#use.user_op(); - let operand_index = r#use.find_index(ctx); - if operations.contains(&user) - && (mma_semantic_consumer(ctx, user, operand_index, fact) - || mma_ordinary_use_forwards_fact( - ctx, - facts, - key, - fact, - user, - operand_index, - )) - { - continue; - } - return Err(invalid(format!( - "shared-MMA carrier has an unsupported outgoing use by `{}` operand {operand_index}", - Operation::get_opid(user, ctx) - ))); - } - } - } - Ok(()) -} - -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(()) -} - -fn verify_smem_mma_story(ctx: &Context, all_ops: &[Ptr]) -> Result { - let ids = smem_mma_semantic_ids(); - let facts = build_mma_facts(ctx, all_ops); - let mut overlays = HashSet::new(); - let mut used_overlays = HashSet::new(); - let mut slices = HashSet::new(); - let mut used_slices = HashSet::new(); - let mut fills = HashSet::new(); - let mut used_fills = HashSet::new(); - let mut scale_loads = HashSet::new(); - let mut used_scale_loads = HashSet::new(); - let mut scale_slices = HashSet::new(); - let mut used_scale_slices = HashSet::new(); - let mut a_loads = HashSet::new(); - let mut used_a_loads = HashSet::new(); - let mut b_partitions = HashSet::new(); - let mut used_b_partitions = HashSet::new(); - let mut gemms = 0usize; - - for operation in all_ops { - let opid = Operation::get_opid(*operation, ctx); - if !ids.contains(&opid) { - continue; - } - if operation.deref(ctx).get_parent_block().is_none() { - return Err(invalid(format!( - "shared/MMA operation `{opid}` is not attached to a block" - ))); - } - let operation_ref = operation.deref(ctx); - if opid == CuteSmemTensorOverlayOp::get_opid_static() { - overlays.insert(*operation); - } else if opid == CuteTiledMmaSliceOp::get_opid_static() { - slices.insert(*operation); - } else if opid == CuteFragmentFillOp::get_opid_static() { - fills.insert(*operation); - } else if opid == CuteMmaLoadScalesOp::get_opid_static() { - scale_loads.insert(*operation); - let semantic = CuteMmaLoadScalesOp::wrap(*operation); - let lane = require_mma_lane( - &facts, - operation_ref.get_operand(0), - "cute.mma_load_scales lane", - )?; - require_same_mma_plan(ctx, lane, *operation, "mma scale load")?; - used_slices.insert(lane); - let scale_a = require_smem_overlay_pair( - &facts, - operation_ref.get_operand(1), - operation_ref.get_operand(2), - "cute.mma_load_scales A", - )?; - let scale_b = require_smem_overlay_pair( - &facts, - operation_ref.get_operand(3), - operation_ref.get_operand(4), - "cute.mma_load_scales B", - )?; - if semantic.scale_a_view(ctx) != CuteSmemTensorOverlayOp::wrap(scale_a).view(ctx) - || semantic.scale_b_view(ctx) != CuteSmemTensorOverlayOp::wrap(scale_b).view(ctx) - { - return Err(invalid( - "cute.mma_load_scales typed views do not match its exact shared overlays", - )); - } - used_overlays.extend([scale_a, scale_b]); - } else if opid == CuteFragmentSliceKOp::get_opid_static() { - scale_slices.insert(*operation); - let load = match one_mma_fact( - &facts, - operation_ref.get_operand(0), - "cute.fragment_slice_k stage scales", - )? { - MmaFact::StageScales(load) => load, - _ => { - return Err(invalid( - "cute.fragment_slice_k must consume cute.mma_load_scales", - )); - } - }; - require_same_mma_plan(ctx, load, *operation, "scale K slice")?; - used_scale_loads.insert(load); - } else if opid == CuteMmaLoadAOp::get_opid_static() { - a_loads.insert(*operation); - let semantic = CuteMmaLoadAOp::wrap(*operation); - let lane = - require_mma_lane(&facts, operation_ref.get_operand(0), "cute.mma_load_a lane")?; - require_same_mma_plan(ctx, lane, *operation, "MMA A load")?; - used_slices.insert(lane); - let overlay = require_smem_overlay_pair( - &facts, - operation_ref.get_operand(1), - operation_ref.get_operand(2), - "cute.mma_load_a", - )?; - if semantic.view(ctx) != CuteSmemTensorOverlayOp::wrap(overlay).view(ctx) { - return Err(invalid( - "cute.mma_load_a typed view does not match its exact shared overlay", - )); - } - used_overlays.insert(overlay); - } else if opid == CuteMmaPartitionBOp::get_opid_static() { - b_partitions.insert(*operation); - let semantic = CuteMmaPartitionBOp::wrap(*operation); - let overlay = require_smem_overlay_pair( - &facts, - operation_ref.get_operand(0), - operation_ref.get_operand(1), - "cute.mma_partition_b", - )?; - if semantic.view(ctx) != CuteSmemTensorOverlayOp::wrap(overlay).view(ctx) { - return Err(invalid( - "cute.mma_partition_b typed view does not match its exact shared overlay", - )); - } - used_overlays.insert(overlay); - } else if opid == CuteTiledGemmOp::get_opid_static() { - gemms += 1; - let semantic = CuteTiledGemmOp::wrap(*operation); - let lane = - require_mma_lane(&facts, operation_ref.get_operand(0), "cute.tiled_gemm lane")?; - let a = match one_mma_fact( - &facts, - operation_ref.get_operand(1), - "cute.tiled_gemm A fragment", - )? { - MmaFact::AFragment(load) => load, - _ => return Err(invalid("cute.tiled_gemm A must come from cute.mma_load_a")), - }; - let b = match ( - one_mma_fact(&facts, operation_ref.get_operand(2), "tiled_gemm B base")?, - one_mma_fact( - &facts, - operation_ref.get_operand(3), - "tiled_gemm B capacity", - )?, - one_mma_fact(&facts, operation_ref.get_operand(4), "tiled_gemm B warp-N")?, - one_mma_fact(&facts, operation_ref.get_operand(5), "tiled_gemm B K-half")?, - ) { - ( - MmaFact::BBase(base), - MmaFact::BCapacity(capacity), - MmaFact::BWarpN(warp_n), - MmaFact::BKHalf(k_half), - ) if base == capacity && base == warp_n && base == k_half => base, - _ => { - return Err(invalid( - "cute.tiled_gemm B carriers must be the four matching results of one mma_partition_b", - )); - } - }; - let scales = match one_mma_fact( - &facts, - operation_ref.get_operand(6), - "cute.tiled_gemm scales", - )? { - MmaFact::SlicedScales(slice) => slice, - _ => { - return Err(invalid( - "cute.tiled_gemm scales must come from cute.fragment_slice_k", - )); - } - }; - let fill = match one_mma_fact( - &facts, - operation_ref.get_operand(7), - "cute.tiled_gemm accumulator", - )? { - MmaFact::Accumulator(fill) => fill, - _ => { - return Err(invalid( - "cute.tiled_gemm accumulator must be the closed recurrence seeded by fragment_fill", - )); - } - }; - for producer in [lane, a, b, scales, fill] { - require_same_mma_plan(ctx, producer, *operation, "tiled GEMM spine")?; - } - let scale_load = match one_mma_fact( - &facts, - scales.deref(ctx).get_operand(0), - "tiled GEMM scale-load origin", - )? { - MmaFact::StageScales(load) => load, - _ => unreachable!("fragment_slice_k was verified above"), - }; - let scale_load_ref = scale_load.deref(ctx); - let a_ref = a.deref(ctx); - let b_ref = b.deref(ctx); - let scale_lane = require_mma_lane( - &facts, - scale_load_ref.get_operand(0), - "mma_load_scales lane", - )?; - let a_lane = require_mma_lane(&facts, a_ref.get_operand(0), "mma_load_a lane")?; - if scale_lane != lane || a_lane != lane { - return Err(invalid( - "tiled GEMM, scale load, and A load must use one tiled-MMA lane slice", - )); - } - require_same_mma_scalar( - &facts, - scale_load_ref.get_operand(5), - a_ref.get_operand(3), - "shared-MMA warp-M", - )?; - require_same_mma_scalar( - &facts, - scale_load_ref.get_operand(6), - b_ref.get_operand(2), - "shared-MMA warp-N", - )?; - require_same_mma_scalar( - &facts, - scales.deref(ctx).get_operand(1), - a_ref.get_operand(4), - "shared-MMA K-half", - )?; - require_same_mma_scalar( - &facts, - a_ref.get_operand(4), - b_ref.get_operand(3), - "shared-MMA K-half", - )?; - if semantic.b_view(ctx) != CuteMmaPartitionBOp::wrap(b).view(ctx) { - return Err(invalid( - "cute.tiled_gemm B view does not match its mma_partition_b producer", - )); - } - used_slices.insert(lane); - used_a_loads.insert(a); - used_b_partitions.insert(b); - used_scale_slices.insert(scales); - used_fills.insert(fill); - } - } - - let found = overlays.len() - + slices.len() - + fills.len() - + scale_loads.len() - + scale_slices.len() - + a_loads.len() - + b_partitions.len() - + gemms; - if found == 0 { - return Ok(()); - } - verify_mma_use_closure(ctx, all_ops, &facts)?; - if gemms == 0 - || overlays != used_overlays - || slices != used_slices - || fills != used_fills - || scale_loads != used_scale_loads - || scale_slices != used_scale_slices - || a_loads != used_a_loads - || b_partitions != used_b_partitions - { - return Err(invalid( - "shared-MMA story is incomplete or has an unconsumed overlay/slice/load/partition/fill producer", - )); - } - Ok(()) -} - -/// 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. -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; - } -} - -fn is_epilogue_pointer_identity_cast(ctx: &Context, operation: Ptr) -> bool { - is_semantic_pointer_carrier_cast(ctx, operation) -} - -fn is_epilogue_writer_plumbing(ctx: &Context, operation: Ptr) -> bool { - let opid = Operation::get_opid(operation, ctx); - opid == MirGotoOp::get_opid_static() - || opid == MirStorageLiveOp::get_opid_static() - || opid == MirStorageDeadOp::get_opid_static() - || opid == MirExtractFieldOp::get_opid_static() - || is_epilogue_pointer_identity_cast(ctx, operation) - || is_closed_local_cell_load(ctx, operation) -} - -/// A non-volatile `mir.load` from a closed compiler-created local cell only -/// re-materializes carrier state (e.g. the epilogue handle struct spilled to -/// an alloca slot); it cannot touch tile memory, so it is admissible plumbing -/// between the writer-protocol anchor operations. -fn is_closed_local_cell_load(ctx: &Context, operation: Ptr) -> bool { - if Operation::get_opid(operation, ctx) != MirLoadOp::get_opid_static() { - return false; - } - let load = MirLoadOp::new(operation); - { - let operation_ref = operation.deref(ctx); - if operation_ref.get_num_operands() != 1 || operation_ref.get_num_results() != 1 { - return false; - } - } - if load.is_volatile(ctx) { - return false; - } - let Some((root, _path)) = local_cell(ctx, load.address_opd(ctx)) else { - return false; - }; - local_cell_is_closed(ctx, root) -} - -/// Accept only the dead immutable reference Rust leaves around a no-inline -/// call on the fieldless, zero-sized TMA store-pipeline token. -fn is_dead_epilogue_store_pipeline_ref(ctx: &Context, operation: Ptr) -> bool { - if Operation::get_opid(operation, ctx) != MirRefOp::get_opid_static() { - return false; - } - let reference = MirRefOp::new(operation); - if reference.verify(ctx).is_err() - || reference - .get_attr_mutable(ctx) - .is_none_or(|mutable| mutable.0) - { - return false; - } - let operation_ref = operation.deref(ctx); - let referent = operation_ref.get_operand(0); - let result = operation_ref.get_result(0); - if result.is_used(ctx) { - return false; - } - - let referent_ty = referent.get_type(ctx); - let referent_ty_ref = referent_ty.deref(ctx); - let Some(store_pipeline) = referent_ty_ref.downcast_ref::() else { - return false; - }; - if store_pipeline.name != "TmaStorePipeline" - || !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 -} - -fn is_epilogue_issuer_plumbing(ctx: &Context, operation: Ptr) -> bool { - let opid = Operation::get_opid(operation, ctx); - is_epilogue_writer_plumbing(ctx, operation) - || opid == MirUndefOp::get_opid_static() - || opid == MirInsertFieldOp::get_opid_static() - || opid == MirConstructStructOp::get_opid_static() - || opid == MirConstructTupleOp::get_opid_static() - || opid == MirConstantOp::get_opid_static() - || opid == MirAddOp::get_opid_static() - || is_dead_epilogue_store_pipeline_ref(ctx, operation) -} - -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", - )) -} - -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(()) -} - -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()) -} - -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(()) -} - -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", - )) - } -} - -#[derive(Clone, Copy)] -struct EpilogueStoreFacts { - operation: Ptr, - source_view: Ptr, - destination_view: Ptr, -} - -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, - }) -} - -fn verify_epilogue_story(ctx: &Context, all_ops: &[Ptr]) -> Result { - let ids = epilogue_semantic_ids(); - let mut overlays = Vec::new(); - let mut slices = Vec::new(); - let mut fragments = Vec::new(); - let mut syncs = Vec::new(); - let mut halves = Vec::new(); - let mut acquires = Vec::new(); - let mut commits = Vec::new(); - let mut tails = Vec::new(); - let mut stores = Vec::new(); - for operation in all_ops { - let opid = Operation::get_opid(*operation, ctx); - if opid == ids[0] { - overlays.push(*operation); - } else if opid == ids[1] { - slices.push(*operation); - } else if opid == ids[2] { - fragments.push(*operation); - } else if opid == ids[3] { - syncs.push(*operation); - } else if opid == ids[4] { - halves.push(*operation); - } else if opid == ids[5] { - acquires.push(*operation); - } else if opid == ids[6] { - commits.push(*operation); - } else if opid == ids[7] { - tails.push(*operation); - } else if opid == ids[8] { - stores.push(*operation); - } - } - - let found = overlays.len() - + slices.len() - + fragments.len() - + syncs.len() - + halves.len() - + acquires.len() - + commits.len() - + tails.len() - + stores.len(); - if found == 0 { - return Ok(()); - } - if overlays.len() != 1 - || slices.len() != 1 - || fragments.len() != 1 - || syncs.len() != 2 - || halves.len() != 2 - || acquires.len() != 1 - || commits.len() != 1 - || tails.len() != 1 - || stores.len() != 2 - { - return Err(invalid(format!( - "epilogue story needs exactly 1 overlay, 1 warp slice, 1 fragment store, 2 syncs, 2 halves, 1 acquire, 2 TMA stores, 1 commit, and 1 tail; found {}, {}, {}, {}, {}, {}, {}, {}, {}", - overlays.len(), - slices.len(), - fragments.len(), - syncs.len(), - halves.len(), - acquires.len(), - stores.len(), - commits.len(), - tails.len() - ))); - } - - let overlay = overlays[0]; - let warp_slice = slices[0]; - let store_fragment = fragments[0]; - let acquire = acquires[0]; - let commit = commits[0]; - let tail = tails[0]; - let overlay_view = CuteEpilogueSmemOverlayOp::wrap(overlay); - let warp_view = CuteEpilogueWarpSliceOp::wrap(warp_slice); - let fragment = CuteEpilogueStoreFragmentOp::wrap(store_fragment); - let overlay_tile = overlay_view - .tile(ctx) - .expect("locally verified epilogue overlay has a tile"); - if warp_view.tile(ctx).as_ref() != Some(&overlay_tile) - || fragment.tile(ctx).as_ref() != Some(&overlay_tile) - { - return Err(invalid( - "epilogue overlay, warp slice, and fragment store must carry the same tile plan", - )); - } - let mma_facts = build_mma_facts(ctx, all_ops); - let accumulator_seed = match one_mma_fact( - &mma_facts, - fragment.accumulator(ctx), - "epilogue fragment accumulator", - )? { - MmaFact::Accumulator(seed) => seed, - _ => { - return Err(invalid( - "epilogue fragment accumulator must be the closed tiled-GEMM recurrence", - )); - } - }; - if CuteFragmentFillOp::wrap(accumulator_seed).plan(ctx) - != Some(overlay_tile.plan.tiled_mma.clone()) - { - return Err(invalid( - "epilogue accumulator plan differs from its shared-MMA story", - )); - } - require_same_epilogue_carrier( - ctx, - warp_view.input_base(ctx), - overlay_view.base(ctx), - "epilogue warp slice must use the base produced by the epilogue overlay", - )?; - require_same_epilogue_carrier( - ctx, - fragment.base(ctx), - warp_view.base(ctx), - "epilogue fragment store must use the base produced by the warp slice", - )?; - require_same_epilogue_carrier( - ctx, - fragment.warp_id(ctx), - warp_view.warp_id(ctx), - "epilogue fragment store must use the warp ID produced by the warp slice", - )?; - require_same_epilogue_carrier( - ctx, - fragment.lane(ctx), - warp_view.lane(ctx), - "epilogue fragment store must use the lane produced by the warp slice", - )?; - - let reusable_sync = syncs - .iter() - .copied() - .find(|operation| { - CuteEpilogueSyncOp::wrap(*operation).phase(ctx) - == Some(CuteEpilogueSyncPhaseAttr::Reusable) - }) - .ok_or_else(|| invalid("epilogue story is missing the Reusable sync"))?; - let ready_sync = syncs - .iter() - .copied() - .find(|operation| { - CuteEpilogueSyncOp::wrap(*operation).phase(ctx) - == Some(CuteEpilogueSyncPhaseAttr::ReadyForTma) - }) - .ok_or_else(|| invalid("epilogue story is missing the ReadyForTma sync"))?; - if reusable_sync == ready_sync { - return Err(invalid( - "epilogue story needs one distinct sync for each phase", - )); - } - for sync in [reusable_sync, ready_sync] { - let sync = CuteEpilogueSyncOp::wrap(sync); - if sync.tile(ctx).as_ref() != Some(&overlay_tile) { - return Err(invalid( - "both epilogue syncs must carry the overlay tile plan", - )); - } - require_same_epilogue_carrier( - ctx, - sync.base(ctx), - overlay_view.base(ctx), - "both epilogue syncs must use the base produced by the epilogue overlay", - )?; - } - - let writer_range = require_linear_protocol_order( - ctx, - "epilogue writer protocol", - &[reusable_sync, store_fragment, ready_sync], - )?; - for operation in &writer_range { - let expected = [reusable_sync, store_fragment, ready_sync].contains(operation); - if !expected && !is_epilogue_writer_plumbing(ctx, *operation) { - return Err(invalid(format!( - "epilogue writer protocol has interposed non-carrier operation `{}`", - Operation::get_opid(*operation, ctx) - ))); - } - } - let writer_events: Vec<_> = writer_range - .iter() - .copied() - .filter(|operation| { - let opid = Operation::get_opid(*operation, ctx); - opid.dialect.to_string() == crate::CUTE_DIALECT_NAME - || opid == MirCallOp::get_opid_static() - }) - .collect(); - if writer_events != [reusable_sync, store_fragment, ready_sync] { - return Err(invalid( - "epilogue writer protocol must be exactly Reusable -> fragment store -> ReadyForTma, whose semantic boundary includes publication to the TMA proxy", - )); - } - - let half_zero = halves - .iter() - .copied() - .find(|operation| { - CuteEpilogueHalfOp::wrap(*operation) - .half(ctx) - .is_some_and(|half| half.0 == 0) - }) - .ok_or_else(|| invalid("epilogue story is missing half 0"))?; - let half_one = halves - .iter() - .copied() - .find(|operation| { - CuteEpilogueHalfOp::wrap(*operation) - .half(ctx) - .is_some_and(|half| half.0 == 1) - }) - .ok_or_else(|| invalid("epilogue story is missing half 1"))?; - if half_zero == half_one { - return Err(invalid("epilogue story needs two distinct shared halves")); - } - for half in [half_zero, half_one] { - let half = CuteEpilogueHalfOp::wrap(half); - if half.tile(ctx).as_ref() != Some(&overlay_tile) { - return Err(invalid( - "both epilogue halves must carry the overlay tile plan", - )); - } - require_same_epilogue_carrier( - ctx, - half.full_base(ctx), - overlay_view.base(ctx), - "both epilogue halves must use the base produced by the epilogue overlay", - )?; - } - - let store_zero = epilogue_store_facts(ctx, stores[0], all_ops)?; - let store_one = epilogue_store_facts(ctx, stores[1], all_ops)?; - for (half, store) in [(half_zero, store_zero), (half_one, store_one)] { - let half = CuteEpilogueHalfOp::wrap(half); - let source = CuteTmaSmemViewOp::wrap(store.source_view); - require_same_epilogue_carrier( - ctx, - source.base(ctx), - half.half_base(ctx), - "each semantic TMA store source base must come from the matching epilogue half through only the compiler's SSA carrier adapter", - )?; - require_same_epilogue_carrier( - ctx, - source.capacity(ctx), - half.capacity(ctx), - "each semantic TMA store source capacity must come from the matching epilogue half through only the compiler's SSA carrier adapter", - )?; - } - if store_zero.operation.deref(ctx).get_operand(1).get_type(ctx) - != store_one.operation.deref(ctx).get_operand(1).get_type(ctx) - { - return Err(invalid( - "both epilogue stores must target the same typed output tile", - )); - } - let destination_zero = CuteTmaGmemViewOp::wrap(store_zero.destination_view); - let destination_one = CuteTmaGmemViewOp::wrap(store_one.destination_view); - require_same_epilogue_carrier( - ctx, - destination_one.descriptor(ctx), - destination_zero.descriptor(ctx), - "both epilogue halves must target the same TMA destination descriptor", - )?; - let semantic_zero = CuteTmaStore2dSemanticOp::wrap(store_zero.operation); - let semantic_one = CuteTmaStore2dSemanticOp::wrap(store_one.operation); - require_same_epilogue_carrier( - ctx, - semantic_one.tile_row(ctx), - semantic_zero.tile_row(ctx), - "both epilogue halves must target the same output tile row", - )?; - require_next_epilogue_column( - ctx, - semantic_zero.tile_column(ctx), - semantic_one.tile_column(ctx), - )?; - - let issuer_range = require_linear_protocol_order( - ctx, - "epilogue TMA issuer protocol", - &[ - half_zero, - store_zero.operation, - half_one, - store_one.operation, - commit, - ], - )?; - for operation in issuer_range { - let opid = Operation::get_opid(operation, ctx); - let allowed_tma_view = opid == CuteTmaGmemViewOp::get_opid_static() - || opid == CuteTmaSmemViewOp::get_opid_static(); - let expected = [ - half_zero, - store_zero.operation, - half_one, - store_one.operation, - commit, - ] - .contains(&operation); - if !allowed_tma_view && !expected && !is_epilogue_issuer_plumbing(ctx, operation) { - return Err(invalid(format!( - "epilogue TMA issuer protocol has interposed operation `{opid}`" - ))); - } - } - - let acquire_pipeline = CuteTmaStoreAcquireOp::wrap(acquire).pipeline(ctx); - let commit_pipeline = CuteTmaStoreCommitOp::wrap(commit).pipeline(ctx); - let tail_pipeline = CuteTmaStoreTailOp::wrap(tail).pipeline(ctx); - if acquire_pipeline != commit_pipeline || acquire_pipeline != tail_pipeline { - return Err(invalid( - "epilogue acquire, commit, and tail must carry the same TMA store-pipeline facts", - )); - } - - require_epilogue_cfg_links( - ctx, - acquire, - reusable_sync, - ready_sync, - half_zero, - commit, - tail, - ) -} - -/// 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::sm100_ops::verify_story(ctx, &scope).map_err(invalid)?; - crate::sm100_attention_ops::verify_story(ctx, &scope).map_err(invalid)?; - } - Ok(()) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::attributes::{ - CutePipelineStateAttr, CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, - CuteTensorLayoutAttr, - }; - use crate::layout::{ComposedLayout, OffsetUnit}; - use dialect_mir::ops::{MirReturnOp, MirUndefOp}; - use dialect_mir::types::{MirPtrType, address_space}; - use pliron::basic_block::BasicBlock; - use pliron::builtin::attributes::TypeAttr; - use pliron::builtin::op_interfaces::{OperandSegmentInterface, SymbolOpInterface}; - use pliron::builtin::ops::ModuleOp; - use pliron::builtin::types::{FP32Type, FunctionType, IntegerType, Signedness}; - - 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) - } - - 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) - } - - 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() - } - - 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) - } - - 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); - } - - 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) - } - - 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 - } - - 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); - } - - 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); - } - - fn append_complete_scheduler_function( - ctx: &mut Context, - module_block: Ptr, - name: &str, - grid: CuteTileGridAttr, - misroute_work_edge: bool, - ) { - let function_type = FunctionType::get(ctx, vec![], vec![]); - let function = Operation::new( - ctx, - MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let function_op = MirFuncOp::new(ctx, function, TypeAttr::new(function_type.into())); - function_op.set_symbol_name(ctx, name.try_into().unwrap()); - function.insert_at_back(module_block, ctx); - let region = function.deref(ctx).get_region(0); - let u64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); - let entry = BasicBlock::new(ctx, None, vec![]); - entry.insert_at_back(region, ctx); - let header = BasicBlock::new(ctx, None, vec![u64_ty, u64_ty]); - header.insert_at_back(region, ctx); - let body = BasicBlock::new(ctx, None, vec![]); - body.insert_at_back(region, ctx); - let exit = BasicBlock::new(ctx, None, vec![]); - exit.insert_at_back(region, ctx); - - let start = CuteSchedulerNew1dOp::new(ctx, grid).get_operation(); - start.insert_at_back(entry, ctx); - let start = CuteSchedulerNew1dOp::wrap(start); - goto( - ctx, - entry, - header, - vec![start.current(ctx), start.stride(ctx)], - ); - - let current = header.deref(ctx).get_argument(0); - let stride = header.deref(ctx).get_argument(1); - let has_work = CuteSchedulerHasWorkOp::new(ctx, current, grid).get_operation(); - has_work.insert_at_back(header, ctx); - let not = Operation::new( - ctx, - MirNotOp::get_concrete_op_info(), - vec![ - CuteSchedulerHasWorkOp::wrap(has_work) - .has_work(ctx) - .get_type(ctx), - ], - vec![CuteSchedulerHasWorkOp::wrap(has_work).has_work(ctx)], - vec![], - 0, - ); - not.insert_at_back(header, ctx); - let finished = not.deref(ctx).get_result(0); - let (true_target, false_target) = if misroute_work_edge { - (body, exit) - } else { - (exit, body) - }; - cond_branch(ctx, header, finished, true_target, false_target); - - let selected = CuteSchedulerCurrentOp::new(ctx, current, grid).get_operation(); - selected.insert_at_back(body, ctx); - CuteWorkTileCoordinatesOp::new(ctx, CuteSchedulerCurrentOp::wrap(selected).work_tile(ctx)) - .get_operation() - .insert_at_back(body, ctx); - let advance = CuteSchedulerAdvanceOp::new(ctx, current, stride, grid).get_operation(); - advance.insert_at_back(body, ctx); - goto( - ctx, - body, - header, - vec![CuteSchedulerAdvanceOp::wrap(advance).next(ctx), stride], - ); - - Operation::new( - ctx, - MirReturnOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ) - .insert_at_back(exit, ctx); - } - - #[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}"); - } - - #[test] - fn accepts_independent_scheduler_stories_in_two_functions_without_mutation() { - let mut ctx = Context::new(); - let (module, block) = module_top(&mut ctx); - append_complete_scheduler_function( - &mut ctx, - block, - "first", - CuteTileGridAttr::new(4, 4, 1), - false, - ); - append_complete_scheduler_function( - &mut ctx, - block, - "second", - CuteTileGridAttr::new(5, 4, 1), - false, - ); - let mut before = Vec::new(); - collect_ops(&ctx, module, &mut before); - - verify_cute_semantics(&ctx, module).unwrap(); - - let mut after = Vec::new(); - collect_ops(&ctx, module, &mut after); - assert_eq!(after, before); - } - - #[test] - fn rejects_a_scheduler_body_on_the_exit_edge_even_when_preorder_looks_ordered() { - let mut ctx = Context::new(); - let (module, block) = module_top(&mut ctx); - append_complete_scheduler_function( - &mut ctx, - block, - "misrouted", - CuteTileGridAttr::new(4, 4, 1), - true, - ); - - let error = verify_cute_semantics(&ctx, module).unwrap_err().to_string(); - assert!(error.contains("scheduler CFG"), "{error}"); - } - - #[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 - ); - } - - #[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}"); - } - - #[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::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::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::local_cell(&ctx, escalated_alias).is_none()); - assert!(!local_cell_is_closed(&ctx, pointer)); - } - - #[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 - )); - } - - #[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}"); - } - - #[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}"); - } - - #[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}" - ); - } - - #[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}"); - } - - #[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}"); - } - - #[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}"); - } - - #[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/cuteir/dialect-cute/src/verify/cfg.rs b/cuda-oxide/cuteir/dialect-cute/src/verify/cfg.rs new file mode 100644 index 0000000000..6a4da1be59 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/verify/cfg.rs @@ -0,0 +1,286 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Function scopes, reachability, dominance, and source-order protocol checks. + +use std::collections::{HashMap, HashSet}; + +use dialect_mir::ops::{MirGotoOp, function::MirFuncOp}; +use pliron::{ + basic_block::BasicBlock, + context::{Context, Ptr}, + linked_list::ContainsLinkedList, + op::Op, + operation::Operation, +}; + +use super::{Result, invalid}; + +/// The semantic verifier treats every MIR function as an independent story +/// namespace. Tests and tools may also place operations directly below a +/// module, in which case the verifier root is their namespace. +fn semantic_scope( + ctx: &Context, + root: Ptr, + operation: Ptr, +) -> Ptr { + let mut current = operation; + while let Some(parent) = current.deref(ctx).get_parent_op(ctx) { + if Operation::get_opid(parent, ctx) == MirFuncOp::get_opid_static() { + return parent; + } + if parent == root { + return root; + } + current = parent; + } + root +} + +/// Partition operations by their enclosing MIR function, keeping standalone IR in the root scope. +pub(super) fn operations_by_scope( + ctx: &Context, + root: Ptr, + all_ops: &[Ptr], +) -> Vec>> { + let mut scopes: HashMap, Vec>> = HashMap::new(); + for operation in all_ops { + let scope = semantic_scope(ctx, root, *operation); + scopes.entry(scope).or_default().push(*operation); + } + scopes.into_values().collect() +} + +/// Test whether any successor path in the same region reaches the destination block. +pub(super) fn block_reaches( + ctx: &Context, + start: Ptr, + destination: Ptr, +) -> bool { + let Some(region) = start.deref(ctx).get_parent_region() else { + return false; + }; + if destination.deref(ctx).get_parent_region() != Some(region) { + return false; + } + let mut worklist = vec![start]; + let mut visited = HashSet::new(); + while let Some(block) = worklist.pop() { + if !visited.insert(block) { + continue; + } + if block == destination { + return true; + } + let Some(terminator) = block.deref(ctx).get_terminator(ctx) else { + continue; + }; + worklist.extend(terminator.deref(ctx).successors()); + } + false +} + +/// Locate an operation's parent block and source-order index, or return `None` if detached. +fn operation_position( + ctx: &Context, + operation: Ptr, +) -> Option<(Ptr, usize)> { + let block = operation.deref(ctx).get_parent_block()?; + let index = block + .deref(ctx) + .iter(ctx) + .position(|candidate| candidate == operation)?; + Some((block, index)) +} + +/// Existential CFG reachability with source order inside one block. This is +/// used only to connect role-specific semantic milestones; each role's local +/// sequence is checked separately by `require_linear_protocol_order` where the +/// importer promises a straight-line call chain. +pub(super) fn operation_reaches( + ctx: &Context, + source: Ptr, + destination: Ptr, +) -> bool { + let Some((source_block, source_index)) = operation_position(ctx, source) else { + return false; + }; + let Some((destination_block, destination_index)) = operation_position(ctx, destination) else { + return false; + }; + if source_block == destination_block { + return source_index < destination_index; + } + let Some(terminator) = source_block.deref(ctx).get_terminator(ctx) else { + return false; + }; + terminator + .deref(ctx) + .successors() + .any(|successor| block_reaches(ctx, successor, destination_block)) +} + +/// Require source order in one block or dominance from the region entry across blocks. +pub(super) fn operation_dominates( + ctx: &Context, + source: Ptr, + destination: Ptr, +) -> bool { + let Some((source_block, source_index)) = operation_position(ctx, source) else { + return false; + }; + let Some((destination_block, destination_index)) = operation_position(ctx, destination) else { + return false; + }; + if source_block == destination_block { + return source_index < destination_index; + } + let Some(region) = source_block.deref(ctx).get_parent_region() else { + return false; + }; + if destination_block.deref(ctx).get_parent_region() != Some(region) { + return false; + } + let blocks: Vec<_> = region.deref(ctx).iter(ctx).collect(); + let Some(entry) = blocks.first().copied() else { + return false; + }; + let all: HashSet<_> = blocks.iter().copied().collect(); + let mut dominators: HashMap, HashSet>> = blocks + .iter() + .copied() + .map(|block| { + if block == entry { + (block, HashSet::from([entry])) + } else { + (block, all.clone()) + } + }) + .collect(); + + loop { + let mut changed = false; + for block in blocks.iter().copied().filter(|block| *block != entry) { + let predecessors: Vec<_> = block + .uses(ctx) + .into_iter() + .filter_map(|edge| edge.user_op().deref(ctx).get_parent_block()) + .collect(); + let mut next = if let Some(first) = predecessors.first() { + dominators.get(first).cloned().unwrap_or_default() + } else { + HashSet::new() + }; + for predecessor in predecessors.iter().skip(1) { + let predecessor_dominators = + dominators.get(predecessor).cloned().unwrap_or_default(); + next.retain(|candidate| predecessor_dominators.contains(candidate)); + } + next.insert(block); + if dominators.get(&block) != Some(&next) { + dominators.insert(block, next); + changed = true; + } + } + if !changed { + break; + } + } + dominators + .get(&destination_block) + .is_some_and(|set| set.contains(&source_block)) +} + +/// Follow one straight-line source sequence across the block splits emitted +/// for Rust call terminators. Every crossed edge must be an unconditional +/// goto into a block with no side entry. +pub(super) fn require_linear_protocol_order( + ctx: &Context, + label: &str, + ordered: &[Ptr], +) -> Result>> { + let Some(first) = ordered.first() else { + return Err(invalid(format!("{label} has no operations to order"))); + }; + let first_block = first + .deref(ctx) + .get_parent_block() + .ok_or_else(|| invalid(format!("{label} contains an operation outside a block")))?; + let region = first_block + .deref(ctx) + .get_parent_region() + .ok_or_else(|| invalid(format!("{label} starts in a block outside a region")))?; + if ordered.iter().any(|operation| { + operation + .deref(ctx) + .get_parent_block() + .and_then(|block| block.deref(ctx).get_parent_region()) + != Some(region) + }) { + return Err(invalid(format!( + "{label} operations must stay in one region" + ))); + } + + let mut current = first_block; + let mut visited = HashSet::new(); + let mut next_expected = 0usize; + let mut range = Vec::new(); + loop { + if !visited.insert(current) { + return Err(invalid(format!( + "{label} loops before reaching its final operation" + ))); + } + let operations: Vec<_> = current.deref(ctx).iter(ctx).collect(); + let start = if current == first_block { + operations + .iter() + .position(|operation| operation == first) + .ok_or_else(|| { + invalid(format!( + "{label} contains an operation missing from its parent block" + )) + })? + } else { + 0 + }; + for operation in operations.into_iter().skip(start) { + range.push(operation); + if ordered.get(next_expected) == Some(&operation) { + next_expected += 1; + if next_expected == ordered.len() { + return Ok(range); + } + } else if ordered.contains(&operation) { + return Err(invalid(format!( + "{label} is not in the required source order" + ))); + } + } + + let terminator = current + .deref(ctx) + .get_terminator(ctx) + .ok_or_else(|| invalid(format!("{label} ends a block before its final operation")))?; + if Operation::get_opid(terminator, ctx) != MirGotoOp::get_opid_static() + || terminator.deref(ctx).get_num_successors() != 1 + { + return Err(invalid(format!( + "{label} crosses an extra exit instead of one unconditional successor" + ))); + } + let successor = terminator.deref(ctx).get_successor(0); + if successor.deref(ctx).get_parent_region() != Some(region) { + return Err(invalid(format!("{label} leaves its source region"))); + } + if successor.preds(ctx).len() != 1 { + return Err(invalid(format!( + "{label} crosses a block with a side entry" + ))); + } + current = successor; + } +} diff --git a/cuda-oxide/cuteir/dialect-cute/src/verify/epilogue.rs b/cuda-oxide/cuteir/dialect-cute/src/verify/epilogue.rs new file mode 100644 index 0000000000..a649af8f8a --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/verify/epilogue.rs @@ -0,0 +1,691 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared output placement and asynchronous store synchronization contracts. + +use dialect_mir::{ + ops::{ + MirAddOp, MirCallOp, MirCastOp, MirConstantOp, MirConstructStructOp, MirConstructTupleOp, + MirExtractFieldOp, MirGotoOp, MirInsertFieldOp, MirLoadOp, MirRefOp, MirStorageDeadOp, + MirStorageLiveOp, MirUndefOp, + }, + types::{MirPtrType, MirStructType, address_space}, +}; +use pliron::{ + builtin::types::IntegerType, + common_traits::Verify, + context::{Context, Ptr}, + op::Op, + operation::Operation, + r#type::Typed, + value::Value, +}; + +use super::{ + Result, + cfg::{operation_reaches, require_linear_protocol_order}, + ghosts::epilogue_semantic_ids, + invalid, + mma::{MmaFact, build_mma_facts, one_mma_fact}, + provenance::{is_semantic_pointer_carrier_cast, local_cell, local_cell_is_closed}, + tma::direct_tma_view_producer, +}; +use crate::{ + attributes::CuteEpilogueSyncPhaseAttr, + epilogue_ops::{ + CuteEpilogueHalfOp, CuteEpilogueSmemOverlayOp, CuteEpilogueStoreFragmentOp, + CuteEpilogueSyncOp, CuteEpilogueWarpSliceOp, CuteTmaStore2dSemanticOp, + CuteTmaStoreAcquireOp, CuteTmaStoreCommitOp, CuteTmaStoreTailOp, + }, + gemm_tma_ops::{CuteTmaGmemViewOp, CuteTmaSmemViewOp}, + smem_mma_ops::CuteFragmentFillOp, +}; + +/// Accept only pointer casts that preserve shared carrier provenance. +fn is_epilogue_pointer_identity_cast(ctx: &Context, operation: Ptr) -> bool { + is_semantic_pointer_carrier_cast(ctx, operation) +} + +/// Recognize carrier-only operations allowed between writer synchronization anchors. +fn is_epilogue_writer_plumbing(ctx: &Context, operation: Ptr) -> bool { + let opid = Operation::get_opid(operation, ctx); + opid == MirGotoOp::get_opid_static() + || opid == MirStorageLiveOp::get_opid_static() + || opid == MirStorageDeadOp::get_opid_static() + || opid == MirExtractFieldOp::get_opid_static() + || is_epilogue_pointer_identity_cast(ctx, operation) + || is_closed_local_cell_load(ctx, operation) +} + +/// A non-volatile `mir.load` from a closed compiler-created local cell only +/// re-materializes carrier state (e.g. the epilogue handle struct spilled to +/// an alloca slot); it cannot touch tile memory, so it is admissible plumbing +/// between the writer-protocol anchor operations. +fn is_closed_local_cell_load(ctx: &Context, operation: Ptr) -> bool { + if Operation::get_opid(operation, ctx) != MirLoadOp::get_opid_static() { + return false; + } + let load = MirLoadOp::new(operation); + { + let operation_ref = operation.deref(ctx); + if operation_ref.get_num_operands() != 1 || operation_ref.get_num_results() != 1 { + return false; + } + } + if load.is_volatile(ctx) { + return false; + } + let Some((root, _path)) = local_cell(ctx, load.address_opd(ctx)) else { + return false; + }; + local_cell_is_closed(ctx, root) +} + +/// Accept only the dead immutable reference Rust leaves around a no-inline +/// call on the fieldless, zero-sized TMA store-pipeline token. +fn is_dead_epilogue_store_pipeline_ref(ctx: &Context, operation: Ptr) -> bool { + if Operation::get_opid(operation, ctx) != MirRefOp::get_opid_static() { + return false; + } + let reference = MirRefOp::new(operation); + if reference.verify(ctx).is_err() + || reference + .get_attr_mutable(ctx) + .is_none_or(|mutable| mutable.0) + { + return false; + } + let operation_ref = operation.deref(ctx); + let referent = operation_ref.get_operand(0); + let result = operation_ref.get_result(0); + if result.is_used(ctx) { + return false; + } + + let referent_ty = referent.get_type(ctx); + let referent_ty_ref = referent_ty.deref(ctx); + let Some(store_pipeline) = referent_ty_ref.downcast_ref::() else { + return false; + }; + if store_pipeline.name != "TmaStorePipeline" + || !store_pipeline.field_names.is_empty() + || !store_pipeline.field_types.is_empty() + || !store_pipeline.mem_to_decl.is_empty() + || !store_pipeline.field_offsets.is_empty() + || store_pipeline.total_size != 0 + || store_pipeline.abi_align != 1 + { + return false; + } + let result_ty = result.get_type(ctx); + let result_ty_ref = result_ty.deref(ctx); + let Some(pointer) = result_ty_ref.downcast_ref::() else { + return false; + }; + if pointer.pointee != referent_ty + || pointer.is_mutable + || pointer.address_space != address_space::GENERIC + { + return false; + } + + let Some(construct) = referent.defining_op() else { + return false; + }; + if Operation::get_opid(construct, ctx) != MirConstructStructOp::get_opid_static() { + return false; + } + let construct_ref = construct.deref(ctx); + if construct_ref.get_num_operands() != 0 + || construct_ref.get_num_results() != 1 + || construct_ref.get_result(0) != referent + { + return false; + } + let uses = referent.uses(ctx); + uses.len() == 1 && uses[0].user_op() == operation && uses[0].find_index(ctx) == 0 +} + +/// Recognize harmless carrier assembly and indexing between TMA issuer anchors. +fn is_epilogue_issuer_plumbing(ctx: &Context, operation: Ptr) -> bool { + let opid = Operation::get_opid(operation, ctx); + is_epilogue_writer_plumbing(ctx, operation) + || opid == MirUndefOp::get_opid_static() + || opid == MirInsertFieldOp::get_opid_static() + || opid == MirConstructStructOp::get_opid_static() + || opid == MirConstructTupleOp::get_opid_static() + || opid == MirConstantOp::get_opid_static() + || opid == MirAddOp::get_opid_static() + || is_dead_epilogue_store_pipeline_ref(ctx, operation) +} + +/// Follow constant field insertions to the requested carrier, rejecting overly deep chains. +fn resolve_inserted_epilogue_field( + ctx: &Context, + mut aggregate: Value, + wanted_index: u32, +) -> Result> { + for _ in 0..16 { + let Some(definer) = aggregate.defining_op() else { + return Ok(None); + }; + if Operation::get_opid(definer, ctx) != MirInsertFieldOp::get_opid_static() { + return Ok(None); + } + let insert = MirInsertFieldOp::new(definer); + let Some(index) = insert.get_attr_insert_index(ctx) else { + return Err(invalid( + "epilogue carrier has an insert_field without a constant index", + )); + }; + let operation = definer.deref(ctx); + if index.0 == wanted_index { + return Ok(Some(operation.get_operand(1))); + } + aggregate = operation.get_operand(0); + } + Err(invalid( + "epilogue carrier insert/extract chain is unexpectedly deep", + )) +} + +/// Peel only lossless SSA carrier plumbing. Memory, arithmetic, control flow, +/// and arbitrary aliases deliberately stop provenance resolution. +fn resolve_epilogue_ssa_carrier(ctx: &Context, mut value: Value) -> Result { + for _ in 0..16 { + let Some(definer) = value.defining_op() else { + return Ok(value); + }; + let opid = Operation::get_opid(definer, ctx); + if opid == MirCastOp::get_opid_static() { + if is_semantic_pointer_carrier_cast(ctx, definer) { + value = definer.deref(ctx).get_operand(0); + continue; + } + return Ok(value); + } + if opid == MirExtractFieldOp::get_opid_static() { + let extract = MirExtractFieldOp::new(definer); + let Some(index) = extract.get_attr_index(ctx) else { + return Err(invalid( + "epilogue carrier has an extract_field without a constant index", + )); + }; + let aggregate = definer.deref(ctx).get_operand(0); + if let Some(aggregate_definer) = aggregate.defining_op() { + let aggregate_id = Operation::get_opid(aggregate_definer, ctx); + if aggregate_id == MirConstructStructOp::get_opid_static() + || aggregate_id == MirConstructTupleOp::get_opid_static() + { + let aggregate_operation = aggregate_definer.deref(ctx); + if usize::try_from(index.0) + .ok() + .is_some_and(|index| index < aggregate_operation.get_num_operands()) + { + value = aggregate_operation.get_operand(index.0 as usize); + continue; + } + } + } + let Some(inserted) = resolve_inserted_epilogue_field(ctx, aggregate, index.0)? else { + return Ok(value); + }; + value = inserted; + continue; + } + return Ok(value); + } + Err(invalid( + "epilogue carrier cast/insert/extract chain is unexpectedly deep", + )) +} + +/// Reject carriers that differ after peeling lossless SSA forwarding. +fn require_same_epilogue_carrier( + ctx: &Context, + actual: Value, + expected: Value, + message: &str, +) -> Result { + if resolve_epilogue_ssa_carrier(ctx, actual)? != resolve_epilogue_ssa_carrier(ctx, expected)? { + return Err(invalid(message)); + } + Ok(()) +} + +/// Read a direct unsigned 64-bit MIR literal used for epilogue coordinates. +fn epilogue_u64_constant(ctx: &Context, value: Value) -> Option { + let integer = value.get_type(ctx); + let integer = integer.deref(ctx); + if !integer + .downcast_ref::() + .is_some_and(|integer| integer.width() == 64 && integer.is_unsigned()) + { + return None; + } + let constant = value.defining_op()?; + if Operation::get_opid(constant, ctx) != MirConstantOp::get_opid_static() { + return None; + } + let attribute = MirConstantOp::new(constant).get_attr_value(ctx)?; + (attribute.value().bw() <= 64).then(|| attribute.value().to_u64()) +} + +/// Require the right output half's column to be the left column plus one. +fn require_next_epilogue_column(ctx: &Context, left: Value, right: Value) -> Result { + let left = resolve_epilogue_ssa_carrier(ctx, left)?; + let right = resolve_epilogue_ssa_carrier(ctx, right)?; + let Some(add) = right.defining_op() else { + return Err(invalid( + "epilogue half 1 must target the tile column immediately after half 0", + )); + }; + if Operation::get_opid(add, ctx) != MirAddOp::get_opid_static() { + return Err(invalid("epilogue half 1 must target `half_0_column + 1`")); + } + let add = add.deref(ctx); + let lhs = resolve_epilogue_ssa_carrier(ctx, add.get_operand(0))?; + let rhs = resolve_epilogue_ssa_carrier(ctx, add.get_operand(1))?; + if !((lhs == left && epilogue_u64_constant(ctx, rhs) == Some(1)) + || (rhs == left && epilogue_u64_constant(ctx, lhs) == Some(1))) + { + return Err(invalid("epilogue half 1 must target `half_0_column + 1`")); + } + Ok(()) +} + +/// Require reachable hand-offs from acquire through writers, issuer, commit, and tail. +pub(super) fn require_epilogue_cfg_links( + ctx: &Context, + acquire: Ptr, + reusable_sync: Ptr, + ready_sync: Ptr, + half_zero: Ptr, + commit: Ptr, + tail: Ptr, +) -> Result { + if operation_reaches(ctx, acquire, reusable_sync) + && operation_reaches(ctx, ready_sync, half_zero) + && operation_reaches(ctx, commit, tail) + { + Ok(()) + } else { + Err(invalid( + "epilogue CFG must connect acquire -> writer hand-off -> TMA issuer -> commit -> tail inside one role-anchored story", + )) + } +} + +/// A semantic output copy and the direct views supplying its storage. +#[derive(Clone, Copy)] +struct EpilogueStoreFacts { + /// The semantic shared-to-global copy being checked. + operation: Ptr, + /// The direct shared view producer carrying the output half. + source_view: Ptr, + /// The direct descriptor-backed global view producer. + destination_view: Ptr, +} + +/// Resolve both copy operands to direct TMA views inside the current function scope. +fn epilogue_store_facts( + ctx: &Context, + operation: Ptr, + scope: &[Ptr], +) -> Result { + let store = CuteTmaStore2dSemanticOp::wrap(operation); + let source_view = direct_tma_view_producer( + ctx, + store.source(ctx), + &CuteTmaSmemViewOp::get_opid_static(), + "epilogue store source", + scope, + )?; + let destination_view = direct_tma_view_producer( + ctx, + store.destination(ctx), + &CuteTmaGmemViewOp::get_opid_static(), + "epilogue store destination", + scope, + )?; + Ok(EpilogueStoreFacts { + operation, + source_view, + destination_view, + }) +} + +/// Validate the single two-half output protocol, carrier origins, and role-local ordering. +pub(super) fn verify_epilogue_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = epilogue_semantic_ids(); + let mut overlays = Vec::new(); + let mut slices = Vec::new(); + let mut fragments = Vec::new(); + let mut syncs = Vec::new(); + let mut halves = Vec::new(); + let mut acquires = Vec::new(); + let mut commits = Vec::new(); + let mut tails = Vec::new(); + let mut stores = Vec::new(); + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if opid == ids[0] { + overlays.push(*operation); + } else if opid == ids[1] { + slices.push(*operation); + } else if opid == ids[2] { + fragments.push(*operation); + } else if opid == ids[3] { + syncs.push(*operation); + } else if opid == ids[4] { + halves.push(*operation); + } else if opid == ids[5] { + acquires.push(*operation); + } else if opid == ids[6] { + commits.push(*operation); + } else if opid == ids[7] { + tails.push(*operation); + } else if opid == ids[8] { + stores.push(*operation); + } + } + + let found = overlays.len() + + slices.len() + + fragments.len() + + syncs.len() + + halves.len() + + acquires.len() + + commits.len() + + tails.len() + + stores.len(); + if found == 0 { + return Ok(()); + } + if overlays.len() != 1 + || slices.len() != 1 + || fragments.len() != 1 + || syncs.len() != 2 + || halves.len() != 2 + || acquires.len() != 1 + || commits.len() != 1 + || tails.len() != 1 + || stores.len() != 2 + { + return Err(invalid(format!( + "epilogue story needs exactly 1 overlay, 1 warp slice, 1 fragment store, 2 syncs, 2 halves, 1 acquire, 2 TMA stores, 1 commit, and 1 tail; found {}, {}, {}, {}, {}, {}, {}, {}, {}", + overlays.len(), + slices.len(), + fragments.len(), + syncs.len(), + halves.len(), + acquires.len(), + stores.len(), + commits.len(), + tails.len() + ))); + } + + let overlay = overlays[0]; + let warp_slice = slices[0]; + let store_fragment = fragments[0]; + let acquire = acquires[0]; + let commit = commits[0]; + let tail = tails[0]; + let overlay_view = CuteEpilogueSmemOverlayOp::wrap(overlay); + let warp_view = CuteEpilogueWarpSliceOp::wrap(warp_slice); + let fragment = CuteEpilogueStoreFragmentOp::wrap(store_fragment); + let overlay_tile = overlay_view + .tile(ctx) + .expect("locally verified epilogue overlay has a tile"); + if warp_view.tile(ctx).as_ref() != Some(&overlay_tile) + || fragment.tile(ctx).as_ref() != Some(&overlay_tile) + { + return Err(invalid( + "epilogue overlay, warp slice, and fragment store must carry the same tile plan", + )); + } + let mma_facts = build_mma_facts(ctx, all_ops); + let accumulator_seed = match one_mma_fact( + &mma_facts, + fragment.accumulator(ctx), + "epilogue fragment accumulator", + )? { + MmaFact::Accumulator(seed) => seed, + _ => { + return Err(invalid( + "epilogue fragment accumulator must be the closed tiled-GEMM recurrence", + )); + } + }; + if CuteFragmentFillOp::wrap(accumulator_seed).plan(ctx) + != Some(overlay_tile.plan.tiled_mma.clone()) + { + return Err(invalid( + "epilogue accumulator plan differs from its shared-MMA story", + )); + } + require_same_epilogue_carrier( + ctx, + warp_view.input_base(ctx), + overlay_view.base(ctx), + "epilogue warp slice must use the base produced by the epilogue overlay", + )?; + require_same_epilogue_carrier( + ctx, + fragment.base(ctx), + warp_view.base(ctx), + "epilogue fragment store must use the base produced by the warp slice", + )?; + require_same_epilogue_carrier( + ctx, + fragment.warp_id(ctx), + warp_view.warp_id(ctx), + "epilogue fragment store must use the warp ID produced by the warp slice", + )?; + require_same_epilogue_carrier( + ctx, + fragment.lane(ctx), + warp_view.lane(ctx), + "epilogue fragment store must use the lane produced by the warp slice", + )?; + + let reusable_sync = syncs + .iter() + .copied() + .find(|operation| { + CuteEpilogueSyncOp::wrap(*operation).phase(ctx) + == Some(CuteEpilogueSyncPhaseAttr::Reusable) + }) + .ok_or_else(|| invalid("epilogue story is missing the Reusable sync"))?; + let ready_sync = syncs + .iter() + .copied() + .find(|operation| { + CuteEpilogueSyncOp::wrap(*operation).phase(ctx) + == Some(CuteEpilogueSyncPhaseAttr::ReadyForTma) + }) + .ok_or_else(|| invalid("epilogue story is missing the ReadyForTma sync"))?; + if reusable_sync == ready_sync { + return Err(invalid( + "epilogue story needs one distinct sync for each phase", + )); + } + for sync in [reusable_sync, ready_sync] { + let sync = CuteEpilogueSyncOp::wrap(sync); + if sync.tile(ctx).as_ref() != Some(&overlay_tile) { + return Err(invalid( + "both epilogue syncs must carry the overlay tile plan", + )); + } + require_same_epilogue_carrier( + ctx, + sync.base(ctx), + overlay_view.base(ctx), + "both epilogue syncs must use the base produced by the epilogue overlay", + )?; + } + + let writer_range = require_linear_protocol_order( + ctx, + "epilogue writer protocol", + &[reusable_sync, store_fragment, ready_sync], + )?; + for operation in &writer_range { + let expected = [reusable_sync, store_fragment, ready_sync].contains(operation); + if !expected && !is_epilogue_writer_plumbing(ctx, *operation) { + return Err(invalid(format!( + "epilogue writer protocol has interposed non-carrier operation `{}`", + Operation::get_opid(*operation, ctx) + ))); + } + } + let writer_events: Vec<_> = writer_range + .iter() + .copied() + .filter(|operation| { + let opid = Operation::get_opid(*operation, ctx); + opid.dialect.to_string() == crate::CUTE_DIALECT_NAME + || opid == MirCallOp::get_opid_static() + }) + .collect(); + if writer_events != [reusable_sync, store_fragment, ready_sync] { + return Err(invalid( + "epilogue writer protocol must be exactly Reusable -> fragment store -> ReadyForTma, whose semantic boundary includes publication to the TMA proxy", + )); + } + + let half_zero = halves + .iter() + .copied() + .find(|operation| { + CuteEpilogueHalfOp::wrap(*operation) + .half(ctx) + .is_some_and(|half| half.0 == 0) + }) + .ok_or_else(|| invalid("epilogue story is missing half 0"))?; + let half_one = halves + .iter() + .copied() + .find(|operation| { + CuteEpilogueHalfOp::wrap(*operation) + .half(ctx) + .is_some_and(|half| half.0 == 1) + }) + .ok_or_else(|| invalid("epilogue story is missing half 1"))?; + if half_zero == half_one { + return Err(invalid("epilogue story needs two distinct shared halves")); + } + for half in [half_zero, half_one] { + let half = CuteEpilogueHalfOp::wrap(half); + if half.tile(ctx).as_ref() != Some(&overlay_tile) { + return Err(invalid( + "both epilogue halves must carry the overlay tile plan", + )); + } + require_same_epilogue_carrier( + ctx, + half.full_base(ctx), + overlay_view.base(ctx), + "both epilogue halves must use the base produced by the epilogue overlay", + )?; + } + + let store_zero = epilogue_store_facts(ctx, stores[0], all_ops)?; + let store_one = epilogue_store_facts(ctx, stores[1], all_ops)?; + for (half, store) in [(half_zero, store_zero), (half_one, store_one)] { + let half = CuteEpilogueHalfOp::wrap(half); + let source = CuteTmaSmemViewOp::wrap(store.source_view); + require_same_epilogue_carrier( + ctx, + source.base(ctx), + half.half_base(ctx), + "each semantic TMA store source base must come from the matching epilogue half through only the compiler's SSA carrier adapter", + )?; + require_same_epilogue_carrier( + ctx, + source.capacity(ctx), + half.capacity(ctx), + "each semantic TMA store source capacity must come from the matching epilogue half through only the compiler's SSA carrier adapter", + )?; + } + if store_zero.operation.deref(ctx).get_operand(1).get_type(ctx) + != store_one.operation.deref(ctx).get_operand(1).get_type(ctx) + { + return Err(invalid( + "both epilogue stores must target the same typed output tile", + )); + } + let destination_zero = CuteTmaGmemViewOp::wrap(store_zero.destination_view); + let destination_one = CuteTmaGmemViewOp::wrap(store_one.destination_view); + require_same_epilogue_carrier( + ctx, + destination_one.descriptor(ctx), + destination_zero.descriptor(ctx), + "both epilogue halves must target the same TMA destination descriptor", + )?; + let semantic_zero = CuteTmaStore2dSemanticOp::wrap(store_zero.operation); + let semantic_one = CuteTmaStore2dSemanticOp::wrap(store_one.operation); + require_same_epilogue_carrier( + ctx, + semantic_one.tile_row(ctx), + semantic_zero.tile_row(ctx), + "both epilogue halves must target the same output tile row", + )?; + require_next_epilogue_column( + ctx, + semantic_zero.tile_column(ctx), + semantic_one.tile_column(ctx), + )?; + + let issuer_range = require_linear_protocol_order( + ctx, + "epilogue TMA issuer protocol", + &[ + half_zero, + store_zero.operation, + half_one, + store_one.operation, + commit, + ], + )?; + for operation in issuer_range { + let opid = Operation::get_opid(operation, ctx); + let allowed_tma_view = opid == CuteTmaGmemViewOp::get_opid_static() + || opid == CuteTmaSmemViewOp::get_opid_static(); + let expected = [ + half_zero, + store_zero.operation, + half_one, + store_one.operation, + commit, + ] + .contains(&operation); + if !allowed_tma_view && !expected && !is_epilogue_issuer_plumbing(ctx, operation) { + return Err(invalid(format!( + "epilogue TMA issuer protocol has interposed operation `{opid}`" + ))); + } + } + + let acquire_pipeline = CuteTmaStoreAcquireOp::wrap(acquire).pipeline(ctx); + let commit_pipeline = CuteTmaStoreCommitOp::wrap(commit).pipeline(ctx); + let tail_pipeline = CuteTmaStoreTailOp::wrap(tail).pipeline(ctx); + if acquire_pipeline != commit_pipeline || acquire_pipeline != tail_pipeline { + return Err(invalid( + "epilogue acquire, commit, and tail must carry the same TMA store-pipeline facts", + )); + } + + require_epilogue_cfg_links( + ctx, + acquire, + reusable_sync, + ready_sync, + half_zero, + commit, + tail, + ) +} diff --git a/cuda-oxide/cuteir/dialect-cute/src/verify/gemv.rs b/cuda-oxide/cuteir/dialect-cute/src/verify/gemv.rs new file mode 100644 index 0000000000..f3da38775e --- /dev/null +++ b/cuda-oxide/cuteir/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, + gemv_ops::{ + 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/cuteir/dialect-cute/src/verify/ghosts.rs b/cuda-oxide/cuteir/dialect-cute/src/verify/ghosts.rs new file mode 100644 index 0000000000..cba692162c --- /dev/null +++ b/cuda-oxide/cuteir/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::{ + epilogue_ops::{ + CuteEpilogueHalfOp, CuteEpilogueSmemOverlayOp, CuteEpilogueStoreFragmentOp, + CuteEpilogueSyncOp, CuteEpilogueWarpSliceOp, CuteTmaStore2dSemanticOp, + CuteTmaStoreAcquireOp, CuteTmaStoreCommitOp, CuteTmaStoreTailOp, + }, + gemm_tma_ops::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}, + gemv_ops::{ + CuteDotOp, CuteScaledViewKTileOp, CuteScaledViewLoadOp, CuteScaledViewMakeOp, + CuteScaledViewRowOp, CuteTensorMake2DOp, + }, + 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, + }, + tensor_ops::{ + 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::copy_ops::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::sm100_ops::semantic_ids()); + ids.extend(crate::sm100_attention_ops::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/cuteir/dialect-cute/src/verify/mma.rs b/cuda-oxide/cuteir/dialect-cute/src/verify/mma.rs new file mode 100644 index 0000000000..385846de55 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/verify/mma.rs @@ -0,0 +1,782 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared tensor placement, register carriers, and tiled matrix-instruction contracts. + +use std::collections::HashSet; + +use dialect_mir::ops::{ + MirCastOp, MirConstructStructOp, MirConstructTupleOp, MirExtractFieldOp, MirFieldAddrOp, + MirInsertFieldOp, MirLoadOp, MirStoreOp, control_flow::MirCondBranchOp, +}; +use pliron::{ + context::{Context, Ptr}, + op::{Op, OpId}, + operation::Operation, + value::Value, +}; + +use super::{ + Result, + ghosts::smem_mma_semantic_ids, + invalid, + provenance::{ + CarrierKey, OrdinaryFacts, branch_successor_operands, is_semantic_pointer_carrier_cast, + local_cell, local_cell_is_closed, poison_missing_forwarding_sources, + propagate_ordinary_plumbing, seed_unknown_integer_carriers, + }, +}; +use crate::{ + attributes::CuteTiledMmaPlanAttr, + epilogue_ops::CuteEpilogueStoreFragmentOp, + smem_mma_ops::{ + CuteFragmentFillOp, CuteFragmentSliceKOp, CuteMmaLoadAOp, CuteMmaLoadScalesOp, + CuteMmaPartitionBOp, CuteSmemTensorOverlayOp, CuteTiledGemmOp, CuteTiledMmaSliceOp, + }, +}; + +/// Read the tiled-MMA plan from a recognized compute operation, or return `None`. +fn mma_plan_for_op( + ctx: &Context, + operation: Ptr, + opid: &OpId, +) -> Option { + if *opid == CuteTiledMmaSliceOp::get_opid_static() { + CuteTiledMmaSliceOp::wrap(operation).plan(ctx) + } else if *opid == CuteFragmentFillOp::get_opid_static() { + CuteFragmentFillOp::wrap(operation).plan(ctx) + } else if *opid == CuteMmaLoadScalesOp::get_opid_static() { + CuteMmaLoadScalesOp::wrap(operation).plan(ctx) + } else if *opid == CuteFragmentSliceKOp::get_opid_static() { + CuteFragmentSliceKOp::wrap(operation).plan(ctx) + } else if *opid == CuteMmaLoadAOp::get_opid_static() { + CuteMmaLoadAOp::wrap(operation).plan(ctx) + } else if *opid == CuteMmaPartitionBOp::get_opid_static() { + CuteMmaPartitionBOp::wrap(operation).plan(ctx) + } else if *opid == CuteTiledGemmOp::get_opid_static() { + CuteTiledGemmOp::wrap(operation).plan(ctx) + } else { + None + } +} + +/// One possible semantic origin of an ordinary MMA carrier or aggregate field. +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +pub(super) enum MmaFact { + /// Lane selection produced by a tiled-MMA slice. + Lane(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Shared allocation base produced by a tensor overlay. + SmemBase(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Element capacity paired with a tensor overlay's shared base. + SmemCapacity( + #[doc = "Operation identifying the semantic origin of this carrier."] Ptr, + ), + /// A complete stage of register scales produced by a scale load. + StageScales( + #[doc = "Operation identifying the semantic origin of this carrier."] Ptr, + ), + /// A K=64 scale selection produced from one stage. + SlicedScales( + #[doc = "Operation identifying the semantic origin of this carrier."] Ptr, + ), + /// Packed A registers produced by a semantic shared load. + AFragment(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Shared B base preserved by lazy partitioning. + BBase(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Element capacity paired with a lazy B partition. + BCapacity(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Warp-N coordinate of a lazy B partition. + BWarpN(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// K-half coordinate of a lazy B partition. + BKHalf(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// FP32 accumulator initialized or updated by the same compute chain. + Accumulator( + #[doc = "Operation identifying the semantic origin of this carrier."] Ptr, + ), + /// An ordinary scalar whose exact SSA identity must remain consistent. + Scalar(#[doc = "Original SSA value lacking semantic provenance."] Value), + /// An input without a proven semantic origin; merging it poisons provenance. + Unknown(#[doc = "Original SSA value lacking semantic provenance."] Value), +} + +/// Return the sole carrier fact, rejecting absent or conflicting facts; callers validate its kind. +pub(super) fn one_mma_fact( + facts: &OrdinaryFacts, + value: Value, + what: &str, +) -> Result { + let Some(found) = facts.at(value) else { + return Err(invalid(format!( + "{what} has no shared-MMA provenance after ordinary carrier forwarding" + ))); + }; + if found.len() != 1 { + return Err(invalid(format!( + "{what} merges carriers from different shared-MMA stories" + ))); + } + Ok(*found.iter().next().expect("one fact")) +} + +/// Require the lane carrier to originate from the expected tiled-MMA slice. +fn require_mma_lane( + facts: &OrdinaryFacts, + value: Value, + what: &str, +) -> Result> { + match one_mma_fact(facts, value, what)? { + MmaFact::Lane(root) => Ok(root), + _ => Err(invalid(format!( + "{what} must come from cute.tiled_mma_slice" + ))), + } +} + +/// Require the shared base and capacity to come from the same typed overlay. +pub(super) fn require_smem_overlay_pair( + facts: &OrdinaryFacts, + base: Value, + capacity: Value, + what: &str, +) -> Result> { + let base_root = match one_mma_fact(facts, base, &format!("{what} base"))? { + MmaFact::SmemBase(root) => root, + _ => { + return Err(invalid(format!( + "{what} base must come from cute.smem_tensor_overlay" + ))); + } + }; + let capacity_root = match one_mma_fact(facts, capacity, &format!("{what} capacity"))? { + MmaFact::SmemCapacity(root) => root, + _ => { + return Err(invalid(format!( + "{what} capacity must come from cute.smem_tensor_overlay" + ))); + } + }; + if base_root != capacity_root { + return Err(invalid(format!( + "{what} mixes base and capacity from different shared overlays" + ))); + } + Ok(base_root) +} + +/// Require two scalar carriers to preserve the same original SSA identity. +fn require_same_mma_scalar( + facts: &OrdinaryFacts, + left: Value, + right: Value, + what: &str, +) -> Result { + if left == right { + return Ok(()); + } + match ( + one_mma_fact(facts, left, what)?, + one_mma_fact(facts, right, what)?, + ) { + (MmaFact::Scalar(left), MmaFact::Scalar(right)) if left == right => Ok(()), + _ => Err(invalid(format!( + "{what} must be one identical ordinary SSA carrier" + ))), + } +} + +/// Propagate compute origins through semantic operations and ordinary carriers, then poison unknown merges. +pub(super) fn build_mma_facts( + ctx: &Context, + operations: &[Ptr], +) -> OrdinaryFacts { + let mut facts = OrdinaryFacts::default(); + for operation in operations { + let opid = Operation::get_opid(*operation, ctx); + let operation_ref = operation.deref(ctx); + if opid == CuteSmemTensorOverlayOp::get_opid_static() { + let overlay = CuteSmemTensorOverlayOp::wrap(*operation); + facts.insert(overlay.base(ctx), MmaFact::SmemBase(*operation)); + facts.insert(overlay.capacity(ctx), MmaFact::SmemCapacity(*operation)); + } else if opid == CuteTiledMmaSliceOp::get_opid_static() { + facts.insert( + CuteTiledMmaSliceOp::wrap(*operation).sliced_lane(ctx), + MmaFact::Lane(*operation), + ); + } else if opid == CuteFragmentFillOp::get_opid_static() { + facts.insert( + CuteFragmentFillOp::wrap(*operation).fragment(ctx), + MmaFact::Accumulator(*operation), + ); + } else if opid == CuteMmaLoadScalesOp::get_opid_static() { + facts.insert( + CuteMmaLoadScalesOp::wrap(*operation).scales(ctx), + MmaFact::StageScales(*operation), + ); + } else if opid == CuteFragmentSliceKOp::get_opid_static() { + facts.insert( + CuteFragmentSliceKOp::wrap(*operation).scales(ctx), + MmaFact::SlicedScales(*operation), + ); + } else if opid == CuteMmaLoadAOp::get_opid_static() { + facts.insert( + CuteMmaLoadAOp::wrap(*operation).fragment(ctx), + MmaFact::AFragment(*operation), + ); + } else if opid == CuteMmaPartitionBOp::get_opid_static() { + facts.insert(operation_ref.get_result(0), MmaFact::BBase(*operation)); + facts.insert(operation_ref.get_result(1), MmaFact::BCapacity(*operation)); + facts.insert(operation_ref.get_result(2), MmaFact::BWarpN(*operation)); + facts.insert(operation_ref.get_result(3), MmaFact::BKHalf(*operation)); + } + } + + loop { + let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); + for operation in operations { + if Operation::get_opid(*operation, ctx) != CuteTiledGemmOp::get_opid_static() { + continue; + } + let operation_ref = operation.deref(ctx); + let Some(accumulator_facts) = facts.at(operation_ref.get_operand(7)).cloned() else { + continue; + }; + for accumulator in accumulator_facts { + if let MmaFact::Accumulator(root) = accumulator { + changed |= facts.insert( + CuteTiledGemmOp::wrap(*operation).accumulator(ctx), + MmaFact::Accumulator(root), + ); + } + } + } + if !changed { + break; + } + } + + seed_unknown_integer_carriers(ctx, operations, &mut facts, 32, MmaFact::Scalar); + seed_unknown_integer_carriers(ctx, operations, &mut facts, 64, MmaFact::Scalar); + loop { + let poisoned = + poison_missing_forwarding_sources(ctx, operations, &mut facts, MmaFact::Unknown); + let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); + if !poisoned && !propagated { + break; + } + } + facts +} + +/// Look up an exact semantic origin inside a carrier's aggregate field path. +fn mma_fact_at_path( + facts: &OrdinaryFacts, + value: Value, + path: &[u32], + fact: MmaFact, +) -> bool { + facts + .at_path(value, path) + .is_some_and(|found| found.contains(&fact)) +} + +/// Check whether an operation consumes the fact in its designated semantic operand position. +fn mma_semantic_consumer( + ctx: &Context, + user: Ptr, + operand_index: usize, + fact: MmaFact, +) -> bool { + let opid = Operation::get_opid(user, ctx); + match fact { + MmaFact::Lane(_) => { + operand_index == 0 + && (opid == CuteMmaLoadScalesOp::get_opid_static() + || opid == CuteMmaLoadAOp::get_opid_static() + || opid == CuteTiledGemmOp::get_opid_static()) + } + MmaFact::SmemBase(_) => { + (opid == CuteMmaLoadScalesOp::get_opid_static() && matches!(operand_index, 1 | 3)) + || (opid == CuteMmaLoadAOp::get_opid_static() && operand_index == 1) + || (opid == CuteMmaPartitionBOp::get_opid_static() && operand_index == 0) + } + MmaFact::SmemCapacity(_) => { + (opid == CuteMmaLoadScalesOp::get_opid_static() && matches!(operand_index, 2 | 4)) + || (opid == CuteMmaLoadAOp::get_opid_static() && operand_index == 2) + || (opid == CuteMmaPartitionBOp::get_opid_static() && operand_index == 1) + } + MmaFact::StageScales(_) => { + opid == CuteFragmentSliceKOp::get_opid_static() && operand_index == 0 + } + MmaFact::SlicedScales(_) => { + opid == CuteTiledGemmOp::get_opid_static() && operand_index == 6 + } + MmaFact::AFragment(_) => opid == CuteTiledGemmOp::get_opid_static() && operand_index == 1, + MmaFact::BBase(_) => opid == CuteTiledGemmOp::get_opid_static() && operand_index == 2, + MmaFact::BCapacity(_) => opid == CuteTiledGemmOp::get_opid_static() && operand_index == 3, + MmaFact::BWarpN(_) => opid == CuteTiledGemmOp::get_opid_static() && operand_index == 4, + MmaFact::BKHalf(_) => opid == CuteTiledGemmOp::get_opid_static() && operand_index == 5, + MmaFact::Accumulator(_) => { + (opid == CuteTiledGemmOp::get_opid_static() && operand_index == 7) + || (opid == CuteEpilogueStoreFragmentOp::get_opid_static() && operand_index == 3) + } + MmaFact::Scalar(_) | MmaFact::Unknown(_) => true, + } +} + +/// Check that a branch forwards the carrier fact to the matching successor argument. +fn mma_branch_forwards_fact( + ctx: &Context, + facts: &OrdinaryFacts, + user: Ptr, + value: Value, + path: &[u32], + fact: MmaFact, +) -> bool { + let operation = user.deref(ctx); + (0..operation.get_num_successors()).any(|successor_index| { + let successor = operation.get_successor(successor_index); + branch_successor_operands(ctx, user, successor_index).is_some_and(|operands| { + operands + .iter() + .copied() + .enumerate() + .any(|(index, incoming)| { + incoming == value + && index < successor.deref(ctx).get_num_arguments() + && mma_fact_at_path( + facts, + successor.deref(ctx).get_argument(index), + path, + fact, + ) + }) + }) + }) +} + +/// Check that ordinary casts, aggregates, or closed local cells forward the tracked fact without escape. +fn mma_ordinary_use_forwards_fact( + ctx: &Context, + facts: &OrdinaryFacts, + key: &CarrierKey, + fact: MmaFact, + user: Ptr, + operand_index: usize, +) -> bool { + let opid = Operation::get_opid(user, ctx); + let operation = user.deref(ctx); + if opid == MirCastOp::get_opid_static() + && operand_index == 0 + && is_semantic_pointer_carrier_cast(ctx, user) + && operation.get_num_results() == 1 + { + return mma_fact_at_path(facts, operation.get_result(0), &key.path, fact); + } + if (opid == MirConstructStructOp::get_opid_static() + || opid == MirConstructTupleOp::get_opid_static()) + && operation.get_num_results() == 1 + { + let mut destination_path = vec![operand_index as u32]; + destination_path.extend_from_slice(&key.path); + return mma_fact_at_path(facts, operation.get_result(0), &destination_path, fact); + } + if opid == MirInsertFieldOp::get_opid_static() + && operation.get_num_operands() == 2 + && operation.get_num_results() == 1 + { + let Some(index) = MirInsertFieldOp::new(user).get_attr_insert_index(ctx) else { + return false; + }; + if operand_index == 0 { + if key.path.first().copied() == Some(index.0) { + return true; + } + return mma_fact_at_path(facts, operation.get_result(0), &key.path, fact); + } + if operand_index == 1 { + let mut destination_path = vec![index.0]; + destination_path.extend_from_slice(&key.path); + return mma_fact_at_path(facts, operation.get_result(0), &destination_path, fact); + } + return false; + } + if opid == MirExtractFieldOp::get_opid_static() + && operand_index == 0 + && operation.get_num_results() == 1 + { + let Some(index) = MirExtractFieldOp::new(user).get_attr_index(ctx) else { + return false; + }; + if key.path.first().copied() != Some(index.0) { + return true; + } + return mma_fact_at_path(facts, operation.get_result(0), &key.path[1..], fact); + } + if opid == MirStoreOp::get_opid_static() && operation.get_num_operands() == 2 { + let store = MirStoreOp::new(user); + if operand_index == 0 { + return local_cell(ctx, store.address_opd(ctx)) + .is_some_and(|(root, _)| local_cell_is_closed(ctx, root)); + } + if operand_index == 1 + && let Some((root, mut destination_path)) = local_cell(ctx, store.address_opd(ctx)) + && local_cell_is_closed(ctx, root) + { + destination_path.extend_from_slice(&key.path); + return mma_fact_at_path(facts, root, &destination_path, fact); + } + return false; + } + if opid == MirLoadOp::get_opid_static() + && operand_index == 0 + && operation.get_num_results() == 1 + { + let load = MirLoadOp::new(user); + let Some((root, cell_path)) = local_cell(ctx, load.address_opd(ctx)) else { + return false; + }; + if !local_cell_is_closed(ctx, root) { + return false; + } + let Some(result_path) = key.path.strip_prefix(cell_path.as_slice()) else { + return true; + }; + return mma_fact_at_path(facts, operation.get_result(0), result_path, fact); + } + if opid == MirFieldAddrOp::get_opid_static() + && operand_index == 0 + && operation.get_num_results() == 1 + && MirFieldAddrOp::new(user) + .get_attr_field_index(ctx) + .is_some() + { + return local_cell(ctx, operation.get_result(0)) + .is_some_and(|(root, _)| local_cell_is_closed(ctx, root)); + } + if opid == MirCondBranchOp::get_opid_static() && operand_index == 0 { + return false; + } + mma_branch_forwards_fact(ctx, facts, user, key.value, &key.path, fact) +} + +/// Reject ordinary uses that discard or expose semantic MMA carriers outside recognized forwarding. +pub(super) fn verify_mma_use_closure( + ctx: &Context, + operations: &[Ptr], + facts: &OrdinaryFacts, +) -> Result { + for (key, found) in &facts.facts { + for fact in found.iter().copied() { + if matches!(fact, MmaFact::Scalar(_) | MmaFact::Unknown(_)) { + continue; + } + for r#use in key.value.uses(ctx) { + let user = r#use.user_op(); + let operand_index = r#use.find_index(ctx); + if operations.contains(&user) + && (mma_semantic_consumer(ctx, user, operand_index, fact) + || mma_ordinary_use_forwards_fact( + ctx, + facts, + key, + fact, + user, + operand_index, + )) + { + continue; + } + return Err(invalid(format!( + "shared-MMA carrier has an unsupported outgoing use by `{}` operand {operand_index}", + Operation::get_opid(user, ctx) + ))); + } + } + } + Ok(()) +} + +/// Require a compute operation's static plan to match the chain's tiled-MMA root. +fn require_same_mma_plan( + ctx: &Context, + left: Ptr, + right: Ptr, + what: &str, +) -> Result { + let left_id = Operation::get_opid(left, ctx); + let right_id = Operation::get_opid(right, ctx); + if mma_plan_for_op(ctx, left, &left_id) != mma_plan_for_op(ctx, right, &right_id) { + return Err(invalid(format!( + "{what} connects shared-MMA operations with different tiled-MMA plans" + ))); + } + Ok(()) +} + +/// Validate shared overlays, lazy B partitions, scale/A fragments, and accumulator provenance together. +pub(super) fn verify_smem_mma_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = smem_mma_semantic_ids(); + let facts = build_mma_facts(ctx, all_ops); + let mut overlays = HashSet::new(); + let mut used_overlays = HashSet::new(); + let mut slices = HashSet::new(); + let mut used_slices = HashSet::new(); + let mut fills = HashSet::new(); + let mut used_fills = HashSet::new(); + let mut scale_loads = HashSet::new(); + let mut used_scale_loads = HashSet::new(); + let mut scale_slices = HashSet::new(); + let mut used_scale_slices = HashSet::new(); + let mut a_loads = HashSet::new(); + let mut used_a_loads = HashSet::new(); + let mut b_partitions = HashSet::new(); + let mut used_b_partitions = HashSet::new(); + let mut gemms = 0usize; + + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if !ids.contains(&opid) { + continue; + } + if operation.deref(ctx).get_parent_block().is_none() { + return Err(invalid(format!( + "shared/MMA operation `{opid}` is not attached to a block" + ))); + } + let operation_ref = operation.deref(ctx); + if opid == CuteSmemTensorOverlayOp::get_opid_static() { + overlays.insert(*operation); + } else if opid == CuteTiledMmaSliceOp::get_opid_static() { + slices.insert(*operation); + } else if opid == CuteFragmentFillOp::get_opid_static() { + fills.insert(*operation); + } else if opid == CuteMmaLoadScalesOp::get_opid_static() { + scale_loads.insert(*operation); + let semantic = CuteMmaLoadScalesOp::wrap(*operation); + let lane = require_mma_lane( + &facts, + operation_ref.get_operand(0), + "cute.mma_load_scales lane", + )?; + require_same_mma_plan(ctx, lane, *operation, "mma scale load")?; + used_slices.insert(lane); + let scale_a = require_smem_overlay_pair( + &facts, + operation_ref.get_operand(1), + operation_ref.get_operand(2), + "cute.mma_load_scales A", + )?; + let scale_b = require_smem_overlay_pair( + &facts, + operation_ref.get_operand(3), + operation_ref.get_operand(4), + "cute.mma_load_scales B", + )?; + if semantic.scale_a_view(ctx) != CuteSmemTensorOverlayOp::wrap(scale_a).view(ctx) + || semantic.scale_b_view(ctx) != CuteSmemTensorOverlayOp::wrap(scale_b).view(ctx) + { + return Err(invalid( + "cute.mma_load_scales typed views do not match its exact shared overlays", + )); + } + used_overlays.extend([scale_a, scale_b]); + } else if opid == CuteFragmentSliceKOp::get_opid_static() { + scale_slices.insert(*operation); + let load = match one_mma_fact( + &facts, + operation_ref.get_operand(0), + "cute.fragment_slice_k stage scales", + )? { + MmaFact::StageScales(load) => load, + _ => { + return Err(invalid( + "cute.fragment_slice_k must consume cute.mma_load_scales", + )); + } + }; + require_same_mma_plan(ctx, load, *operation, "scale K slice")?; + used_scale_loads.insert(load); + } else if opid == CuteMmaLoadAOp::get_opid_static() { + a_loads.insert(*operation); + let semantic = CuteMmaLoadAOp::wrap(*operation); + let lane = + require_mma_lane(&facts, operation_ref.get_operand(0), "cute.mma_load_a lane")?; + require_same_mma_plan(ctx, lane, *operation, "MMA A load")?; + used_slices.insert(lane); + let overlay = require_smem_overlay_pair( + &facts, + operation_ref.get_operand(1), + operation_ref.get_operand(2), + "cute.mma_load_a", + )?; + if semantic.view(ctx) != CuteSmemTensorOverlayOp::wrap(overlay).view(ctx) { + return Err(invalid( + "cute.mma_load_a typed view does not match its exact shared overlay", + )); + } + used_overlays.insert(overlay); + } else if opid == CuteMmaPartitionBOp::get_opid_static() { + b_partitions.insert(*operation); + let semantic = CuteMmaPartitionBOp::wrap(*operation); + let overlay = require_smem_overlay_pair( + &facts, + operation_ref.get_operand(0), + operation_ref.get_operand(1), + "cute.mma_partition_b", + )?; + if semantic.view(ctx) != CuteSmemTensorOverlayOp::wrap(overlay).view(ctx) { + return Err(invalid( + "cute.mma_partition_b typed view does not match its exact shared overlay", + )); + } + used_overlays.insert(overlay); + } else if opid == CuteTiledGemmOp::get_opid_static() { + gemms += 1; + let semantic = CuteTiledGemmOp::wrap(*operation); + let lane = + require_mma_lane(&facts, operation_ref.get_operand(0), "cute.tiled_gemm lane")?; + let a = match one_mma_fact( + &facts, + operation_ref.get_operand(1), + "cute.tiled_gemm A fragment", + )? { + MmaFact::AFragment(load) => load, + _ => return Err(invalid("cute.tiled_gemm A must come from cute.mma_load_a")), + }; + let b = match ( + one_mma_fact(&facts, operation_ref.get_operand(2), "tiled_gemm B base")?, + one_mma_fact( + &facts, + operation_ref.get_operand(3), + "tiled_gemm B capacity", + )?, + one_mma_fact(&facts, operation_ref.get_operand(4), "tiled_gemm B warp-N")?, + one_mma_fact(&facts, operation_ref.get_operand(5), "tiled_gemm B K-half")?, + ) { + ( + MmaFact::BBase(base), + MmaFact::BCapacity(capacity), + MmaFact::BWarpN(warp_n), + MmaFact::BKHalf(k_half), + ) if base == capacity && base == warp_n && base == k_half => base, + _ => { + return Err(invalid( + "cute.tiled_gemm B carriers must be the four matching results of one mma_partition_b", + )); + } + }; + let scales = match one_mma_fact( + &facts, + operation_ref.get_operand(6), + "cute.tiled_gemm scales", + )? { + MmaFact::SlicedScales(slice) => slice, + _ => { + return Err(invalid( + "cute.tiled_gemm scales must come from cute.fragment_slice_k", + )); + } + }; + let fill = match one_mma_fact( + &facts, + operation_ref.get_operand(7), + "cute.tiled_gemm accumulator", + )? { + MmaFact::Accumulator(fill) => fill, + _ => { + return Err(invalid( + "cute.tiled_gemm accumulator must be the closed recurrence seeded by fragment_fill", + )); + } + }; + for producer in [lane, a, b, scales, fill] { + require_same_mma_plan(ctx, producer, *operation, "tiled GEMM spine")?; + } + let scale_load = match one_mma_fact( + &facts, + scales.deref(ctx).get_operand(0), + "tiled GEMM scale-load origin", + )? { + MmaFact::StageScales(load) => load, + _ => unreachable!("fragment_slice_k was verified above"), + }; + let scale_load_ref = scale_load.deref(ctx); + let a_ref = a.deref(ctx); + let b_ref = b.deref(ctx); + let scale_lane = require_mma_lane( + &facts, + scale_load_ref.get_operand(0), + "mma_load_scales lane", + )?; + let a_lane = require_mma_lane(&facts, a_ref.get_operand(0), "mma_load_a lane")?; + if scale_lane != lane || a_lane != lane { + return Err(invalid( + "tiled GEMM, scale load, and A load must use one tiled-MMA lane slice", + )); + } + require_same_mma_scalar( + &facts, + scale_load_ref.get_operand(5), + a_ref.get_operand(3), + "shared-MMA warp-M", + )?; + require_same_mma_scalar( + &facts, + scale_load_ref.get_operand(6), + b_ref.get_operand(2), + "shared-MMA warp-N", + )?; + require_same_mma_scalar( + &facts, + scales.deref(ctx).get_operand(1), + a_ref.get_operand(4), + "shared-MMA K-half", + )?; + require_same_mma_scalar( + &facts, + a_ref.get_operand(4), + b_ref.get_operand(3), + "shared-MMA K-half", + )?; + if semantic.b_view(ctx) != CuteMmaPartitionBOp::wrap(b).view(ctx) { + return Err(invalid( + "cute.tiled_gemm B view does not match its mma_partition_b producer", + )); + } + used_slices.insert(lane); + used_a_loads.insert(a); + used_b_partitions.insert(b); + used_scale_slices.insert(scales); + used_fills.insert(fill); + } + } + + let found = overlays.len() + + slices.len() + + fills.len() + + scale_loads.len() + + scale_slices.len() + + a_loads.len() + + b_partitions.len() + + gemms; + if found == 0 { + return Ok(()); + } + verify_mma_use_closure(ctx, all_ops, &facts)?; + if gemms == 0 + || overlays != used_overlays + || slices != used_slices + || fills != used_fills + || scale_loads != used_scale_loads + || scale_slices != used_scale_slices + || a_loads != used_a_loads + || b_partitions != used_b_partitions + { + return Err(invalid( + "shared-MMA story is incomplete or has an unconsumed overlay/slice/load/partition/fill producer", + )); + } + Ok(()) +} diff --git a/cuda-oxide/cuteir/dialect-cute/src/verify/mod.rs b/cuda-oxide/cuteir/dialect-cute/src/verify/mod.rs new file mode 100644 index 0000000000..d1192d0623 --- /dev/null +++ b/cuda-oxide/cuteir/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::{ + epilogue_ops::{ + CuteEpilogueHalfOp, CuteEpilogueSmemOverlayOp, CuteEpilogueStoreFragmentOp, + CuteEpilogueSyncOp, CuteEpilogueWarpSliceOp, CuteTmaStore2dSemanticOp, + CuteTmaStoreAcquireOp, CuteTmaStoreCommitOp, CuteTmaStoreTailOp, + }, + gemm_tma_ops::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}, + gemv_ops::{ + CuteDotOp, CuteScaledViewKTileOp, CuteScaledViewLoadOp, CuteScaledViewMakeOp, + CuteScaledViewRowOp, CuteTensorMake2DOp, + }, + ops::{ + CuteAssumeDivOp, CuteCopyG2SOp, CuteCopyOp, CuteLdmatrixOp, CuteTmaLoad2dOp, + CuteTmaStore2dOp, + }, + 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, + }, + tensor_ops::{ + CuteTensorBaseOp, CuteTensorIsFullOp, CuteTensorLoadIntoOp, CuteTensorMakeOp, + CuteTensorSliceOp, CuteTensorStoreElementAbsOp, CuteTensorStoreFromOp, + CuteTensorZippedDivideOp, + }, +}; + +/// 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::copy_ops::CuteStoreGlobalNoAllocateOp::get_opid_static() { + Some(verify_as!(crate::copy_ops::CuteStoreGlobalNoAllocateOp)) + } else if let Some(result) = crate::sm100_attention_ops::verify_local(ctx, operation) { + Some(result) + } else if let Some(result) = crate::sm100_ops::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::sm100_ops::verify_story(ctx, &scope).map_err(invalid)?; + crate::sm100_attention_ops::verify_story(ctx, &scope).map_err(invalid)?; + } + Ok(()) +} diff --git a/cuda-oxide/cuteir/dialect-cute/src/verify/pipeline.rs b/cuda-oxide/cuteir/dialect-cute/src/verify/pipeline.rs new file mode 100644 index 0000000000..6bb083f482 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/verify/pipeline.rs @@ -0,0 +1,607 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Barrier-ring configuration and producer/consumer cursor contracts. + +use std::collections::{HashMap, HashSet}; + +use pliron::{ + context::{Context, Ptr}, + op::{Op, OpId}, + operation::Operation, + r#type::Typed, + value::Value, +}; + +use super::{ + Result, + cfg::{operation_dominates, operation_reaches}, + ghosts::pipeline_semantic_ids, + invalid, + provenance::{ + OrdinaryFacts, poison_missing_forwarding_sources, propagate_ordinary_plumbing, + seed_unknown_integer_carriers, + }, +}; +use crate::{ + attributes::{CutePipelineRoleAttr, CutePipelineStateAttr}, + gemm_tma_ops::CuteTmaCopy2dOp, + pipeline_ops::{ + CutePipelineConsumerReleaseOp, CutePipelineConsumerWaitOp, CutePipelineProducerAcquireOp, + CutePipelineProducerExpectTxOp, CutePipelineProducerTailOp, CutePipelineStateAdvanceOp, + CutePipelineStateNewOp, CutePipelineStateSlotOp, CuteTmaLoadPipelineInitOp, + CuteTmaLoadPipelineMakeOp, + }, + types::CuteTmaViewType, +}; + +/// Static ring facts and the constructor identifying one load pipeline. +#[derive(Clone, Copy)] +struct LoadPipelineFacts { + /// Direct constructor that defines the compiler-only pipeline handle. + make: Ptr, + /// Number of full and empty barrier slots in each ring. + stages: u64, + /// Total bytes expected by one full barrier for its stage. + transaction_bytes: u32, +} + +/// Require a direct pipeline constructor and recover its ring configuration. +fn resolve_load_pipeline(ctx: &Context, value: Value) -> Result { + let Some(make) = value.defining_op() else { + return Err(invalid( + "TMA load-pipeline handle reached a block argument; a direct tma_load_pipeline_make result is required", + )); + }; + let producer_id = Operation::get_opid(make, ctx); + if producer_id != CuteTmaLoadPipelineMakeOp::get_opid_static() { + return Err(invalid(format!( + "TMA load-pipeline handle is produced by `{producer_id}` instead of `cute.tma_load_pipeline_make`" + ))); + } + let make_op = CuteTmaLoadPipelineMakeOp::wrap(make); + if make_op.pipeline(ctx) != value { + return Err(invalid( + "TMA load-pipeline operand is not the make operation's pipeline result", + )); + } + let pipeline = make_op + .pipeline_type(ctx) + .ok_or_else(|| invalid("TMA load-pipeline make has no pipeline type"))?; + Ok(LoadPipelineFacts { + make, + stages: pipeline.stages, + transaction_bytes: pipeline.transaction_bytes, + }) +} + +/// Read the producer/consumer role and ring size from a recognized cursor operation. +fn pipeline_state_for_op( + ctx: &Context, + operation: Ptr, + opid: &OpId, +) -> Option { + if *opid == CutePipelineStateNewOp::get_opid_static() { + CutePipelineStateNewOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineStateSlotOp::get_opid_static() { + CutePipelineStateSlotOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineStateAdvanceOp::get_opid_static() { + CutePipelineStateAdvanceOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineProducerAcquireOp::get_opid_static() { + CutePipelineProducerAcquireOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineProducerExpectTxOp::get_opid_static() { + CutePipelineProducerExpectTxOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineConsumerWaitOp::get_opid_static() { + CutePipelineConsumerWaitOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineConsumerReleaseOp::get_opid_static() { + CutePipelineConsumerReleaseOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineProducerTailOp::get_opid_static() { + CutePipelineProducerTailOp::wrap(operation).state(ctx) + } else { + None + } +} + +/// Find a lifecycle operation's pipeline handle; cursor-only operations return `None`. +fn pipeline_operand(ctx: &Context, operation: Ptr, opid: &OpId) -> Option { + if *opid == CuteTmaLoadPipelineInitOp::get_opid_static() { + Some(CuteTmaLoadPipelineInitOp::wrap(operation).pipeline(ctx)) + } else if *opid == CutePipelineProducerAcquireOp::get_opid_static() { + Some(CutePipelineProducerAcquireOp::wrap(operation).pipeline(ctx)) + } else if *opid == CutePipelineProducerExpectTxOp::get_opid_static() { + Some(CutePipelineProducerExpectTxOp::wrap(operation).pipeline(ctx)) + } else if *opid == CutePipelineConsumerWaitOp::get_opid_static() { + Some(CutePipelineConsumerWaitOp::wrap(operation).pipeline(ctx)) + } else if *opid == CutePipelineConsumerReleaseOp::get_opid_static() { + Some(CutePipelineConsumerReleaseOp::wrap(operation).pipeline(ctx)) + } else if *opid == CutePipelineProducerTailOp::get_opid_static() { + Some(CutePipelineProducerTailOp::wrap(operation).pipeline(ctx)) + } else { + None + } +} + +/// Sum the static byte sizes of every direct TMA copy completed by one +/// `expect_tx` result. The result may not escape through any other operand. +fn expected_tma_transaction_bytes(ctx: &Context, operation: Ptr) -> Result { + let barrier = CutePipelineProducerExpectTxOp::wrap(operation).completion_barrier(ctx); + let uses = barrier.uses(ctx); + if uses.is_empty() { + return Err(invalid( + "cute.pipeline_producer_expect_tx needs at least one direct cute.tma_copy_2d user", + )); + } + let mut copies = HashSet::new(); + let mut bytes = 0u64; + for r#use in uses { + let user = r#use.user_op(); + if Operation::get_opid(user, ctx) != CuteTmaCopy2dOp::get_opid_static() + || r#use.find_index(ctx) != 4 + || CuteTmaCopy2dOp::wrap(user).completion_barrier(ctx) != barrier + { + return Err(invalid( + "cute.pipeline_producer_expect_tx result may only be the completion-barrier operand of cute.tma_copy_2d", + )); + } + if !copies.insert(user) { + continue; + } + let copy = CuteTmaCopy2dOp::wrap(user); + let source_ty = copy.source(ctx).get_type(ctx); + let source_ty = source_ty.deref(ctx); + let view = source_ty + .downcast_ref::() + .ok_or_else(|| invalid("TMA copy attached to expect_tx has no TMA view type"))?; + let tile_bytes = view.tile_bytes(ctx).ok_or_else(|| { + invalid("TMA copy attached to expect_tx has no static tile byte count") + })?; + bytes = bytes + .checked_add(tile_bytes) + .ok_or_else(|| invalid("TMA transaction byte total overflows u64"))?; + } + Ok(bytes) +} + +/// Origin of one ordinary cursor component, including untrusted incoming values. +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +enum PipelineFact { + /// Ring slot derived from this cursor-creation operation. + Slot(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Barrier parity derived from this cursor-creation operation. + Phase(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Scalar without a proven cursor origin; it cannot satisfy lifecycle checks. + Unknown(#[doc = "Original SSA value lacking semantic provenance."] Value), +} + +/// Require an unambiguous slot or phase fact rooted at cursor creation. +fn pipeline_state_root( + facts: &OrdinaryFacts, + value: Value, + slot: bool, + what: &str, +) -> Result> { + let Some(found) = facts.at(value) else { + return Err(invalid(format!( + "{what} has no pipeline-state provenance after ordinary SSA forwarding" + ))); + }; + if found.len() != 1 { + return Err(invalid(format!( + "{what} merges values from different pipeline-state stories" + ))); + } + match found.iter().next().copied() { + Some(PipelineFact::Slot(root)) if slot => Ok(root), + Some(PipelineFact::Phase(root)) if !slot => Ok(root), + _ => Err(invalid(format!( + "{what} is not derived from the required cute.pipeline_state_new result" + ))), + } +} + +/// Track slot/phase origins through advances and forwarding, preserving unknown inputs for consumer checks. +fn build_pipeline_facts( + ctx: &Context, + operations: &[Ptr], +) -> OrdinaryFacts { + let mut facts = OrdinaryFacts::default(); + for operation in operations { + if Operation::get_opid(*operation, ctx) == CutePipelineStateNewOp::get_opid_static() { + let state = CutePipelineStateNewOp::wrap(*operation); + facts.insert(state.slot(ctx), PipelineFact::Slot(*operation)); + facts.insert(state.phase(ctx), PipelineFact::Phase(*operation)); + } + } + + loop { + let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); + for operation in operations { + if Operation::get_opid(*operation, ctx) != CutePipelineStateAdvanceOp::get_opid_static() + { + continue; + } + let advance = CutePipelineStateAdvanceOp::wrap(*operation); + let Some(slot_facts) = facts.at(advance.slot(ctx)).cloned() else { + continue; + }; + let Some(phase_facts) = facts.at(advance.phase(ctx)).cloned() else { + continue; + }; + for slot in &slot_facts { + let PipelineFact::Slot(root) = slot else { + continue; + }; + if phase_facts.contains(&PipelineFact::Phase(*root)) + && CutePipelineStateNewOp::wrap(*root).state(ctx) == advance.state(ctx) + { + changed |= facts.insert(advance.next_slot(ctx), PipelineFact::Slot(*root)); + changed |= facts.insert(advance.next_phase(ctx), PipelineFact::Phase(*root)); + } + } + } + if !changed { + break; + } + } + + seed_unknown_integer_carriers(ctx, operations, &mut facts, 32, PipelineFact::Unknown); + loop { + let poisoned = + poison_missing_forwarding_sources(ctx, operations, &mut facts, PipelineFact::Unknown); + let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); + if !poisoned && !propagated { + break; + } + } + facts +} + +/// Static lifecycle sites collected for one producer or consumer cursor root. +#[derive(Default)] +struct PipelineStateEvents { + /// Sites that widen the cursor slot for addressing. + slots: Vec>, + /// Sites that advance this cursor's slot and phase. + advances: Vec>, + /// Producer sites waiting for an empty stage. + acquires: Vec>, + /// Producer sites announcing the stage's transaction bytes. + expects: Vec>, + /// Consumer sites waiting for a full stage. + waits: Vec>, + /// Consumer sites releasing the stage after reading. + releases: Vec>, + /// Producer exit sites draining the empty-barrier ring. + tails: Vec>, +} + +/// Require dominance from cursor creation through acquire/expect/advance and a reachable exit tail. +pub(super) fn require_producer_pipeline_order( + ctx: &Context, + root: Ptr, + acquire: Ptr, + expect: Ptr, + advance: Ptr, + tail: Ptr, +) -> Result { + if operation_dominates(ctx, root, acquire) + && operation_dominates(ctx, acquire, expect) + && operation_dominates(ctx, expect, advance) + && operation_reaches(ctx, advance, tail) + { + Ok(()) + } else { + Err(invalid( + "producer pipeline CFG must order state_new -> acquire -> expect_tx -> advance, with tail reachable only after the producer loop", + )) + } +} + +/// Validate handle ownership, complete cursor lifecycles, CFG order, and expected TMA byte totals. +pub(super) fn verify_pipeline_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = pipeline_semantic_ids(); + let facts = build_pipeline_facts(ctx, all_ops); + let makes: Vec<_> = all_ops + .iter() + .copied() + .filter(|operation| { + Operation::get_opid(*operation, ctx) == CuteTmaLoadPipelineMakeOp::get_opid_static() + }) + .collect(); + let mut init_by_make: HashMap, Vec>> = HashMap::new(); + let mut events: HashMap, PipelineStateEvents> = HashMap::new(); + let mut binding: HashMap, Ptr> = HashMap::new(); + + for make in &makes { + if make.deref(ctx).get_parent_block().is_none() { + return Err(invalid("TMA load-pipeline make is not attached to a block")); + } + let handle = CuteTmaLoadPipelineMakeOp::wrap(*make).pipeline(ctx); + let uses = handle.uses(ctx); + if uses.is_empty() { + return Err(invalid( + "cute.tma_load_pipeline_make has no lifecycle consumer", + )); + } + for r#use in uses { + let user = r#use.user_op(); + let user_id = Operation::get_opid(user, ctx); + if !ids.contains(&user_id) + || r#use.find_index(ctx) != 0 + || pipeline_operand(ctx, user, &user_id) != Some(handle) + { + return Err(invalid(format!( + "TMA load-pipeline handle has an invalid use by `{user_id}`" + ))); + } + } + } + + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if !ids.contains(&opid) { + continue; + } + if operation.deref(ctx).get_parent_block().is_none() { + return Err(invalid(format!( + "pipeline operation `{opid}` is not attached to a block" + ))); + } + if opid == CutePipelineStateNewOp::get_opid_static() { + events.entry(*operation).or_default(); + continue; + } + + let state_root = if opid == CutePipelineStateSlotOp::get_opid_static() { + let semantic = CutePipelineStateSlotOp::wrap(*operation); + let root = pipeline_state_root( + &facts, + semantic.slot(ctx), + true, + "cute.pipeline_state_slot slot", + )?; + if semantic.state(ctx) != CutePipelineStateNewOp::wrap(root).state(ctx) { + return Err(invalid( + "cute.pipeline_state_slot facts differ from its pipeline_state_new root", + )); + } + events.entry(root).or_default().slots.push(*operation); + Some(root) + } else if opid == CutePipelineStateAdvanceOp::get_opid_static() { + let semantic = CutePipelineStateAdvanceOp::wrap(*operation); + let slot_root = pipeline_state_root( + &facts, + semantic.slot(ctx), + true, + "cute.pipeline_state_advance slot", + )?; + let phase_root = pipeline_state_root( + &facts, + semantic.phase(ctx), + false, + "cute.pipeline_state_advance phase", + )?; + if slot_root != phase_root + || semantic.state(ctx) != CutePipelineStateNewOp::wrap(slot_root).state(ctx) + { + return Err(invalid( + "cute.pipeline_state_advance must consume slot and phase from one matching pipeline_state_new", + )); + } + events + .entry(slot_root) + .or_default() + .advances + .push(*operation); + Some(slot_root) + } else if opid == CuteTmaLoadPipelineMakeOp::get_opid_static() { + None + } else if opid == CuteTmaLoadPipelineInitOp::get_opid_static() { + let pipeline = resolve_load_pipeline( + ctx, + CuteTmaLoadPipelineInitOp::wrap(*operation).pipeline(ctx), + )?; + if !makes.contains(&pipeline.make) { + return Err(invalid( + "pipeline init refers to a make outside its semantic scope", + )); + } + init_by_make + .entry(pipeline.make) + .or_default() + .push(*operation); + None + } else { + let state = pipeline_state_for_op(ctx, *operation, &opid) + .ok_or_else(|| invalid(format!("pipeline operation `{opid}` has no state")))?; + let operation_ref = operation.deref(ctx); + let slot_root = pipeline_state_root( + &facts, + operation_ref.get_operand(1), + true, + &format!("`{opid}` slot"), + )?; + let needs_phase = opid == CutePipelineProducerAcquireOp::get_opid_static() + || opid == CutePipelineConsumerWaitOp::get_opid_static() + || opid == CutePipelineProducerTailOp::get_opid_static(); + if needs_phase { + let phase_root = pipeline_state_root( + &facts, + operation_ref.get_operand(2), + false, + &format!("`{opid}` phase"), + )?; + if phase_root != slot_root { + return Err(invalid(format!( + "`{opid}` mixes slot and phase from different pipeline states" + ))); + } + } + if CutePipelineStateNewOp::wrap(slot_root).state(ctx) != Some(state) { + return Err(invalid(format!( + "`{opid}` state facts differ from its pipeline_state_new root" + ))); + } + + let handle = pipeline_operand(ctx, *operation, &opid) + .ok_or_else(|| invalid(format!("`{opid}` has no load-pipeline handle")))?; + let pipeline = resolve_load_pipeline(ctx, handle)?; + if !makes.contains(&pipeline.make) { + return Err(invalid( + "pipeline lifecycle operation refers to a make outside its semantic scope", + )); + } + if state.stages != pipeline.stages { + return Err(invalid(format!( + "pipeline state on `{opid}` has {} stages but its handle has {}", + state.stages, pipeline.stages + ))); + } + if let Some(previous) = binding.insert(slot_root, pipeline.make) + && previous != pipeline.make + { + return Err(invalid( + "one ordinary pipeline state is used with multiple load-pipeline handles", + )); + } + + let root_events = events.entry(slot_root).or_default(); + if opid == CutePipelineProducerAcquireOp::get_opid_static() { + root_events.acquires.push(*operation); + } else if opid == CutePipelineProducerExpectTxOp::get_opid_static() { + root_events.expects.push(*operation); + let copied = expected_tma_transaction_bytes(ctx, *operation)?; + if copied != u64::from(pipeline.transaction_bytes) { + return Err(invalid(format!( + "cute.pipeline_producer_expect_tx promises {} bytes, but its direct TMA copies move {copied} bytes", + pipeline.transaction_bytes + ))); + } + } else if opid == CutePipelineConsumerWaitOp::get_opid_static() { + root_events.waits.push(*operation); + } else if opid == CutePipelineConsumerReleaseOp::get_opid_static() { + root_events.releases.push(*operation); + } else if opid == CutePipelineProducerTailOp::get_opid_static() { + root_events.tails.push(*operation); + } + Some(slot_root) + }; + + if let Some(root) = state_root { + let root_state = CutePipelineStateNewOp::wrap(root) + .state(ctx) + .expect("locally verified pipeline state root"); + if let Some(state) = pipeline_state_for_op(ctx, *operation, &opid) + && state != root_state + { + return Err(invalid(format!( + "`{opid}` does not carry its pipeline-state root's role and stage count" + ))); + } + } + } + + let mut roles_by_make: HashMap, (usize, usize)> = HashMap::new(); + for (root, root_events) in &events { + let state = CutePipelineStateNewOp::wrap(*root) + .state(ctx) + .expect("locally verified pipeline state root"); + let Some(make) = binding.get(root).copied() else { + return Err(invalid( + "cute.pipeline_state_new is not connected to any load-pipeline lifecycle", + )); + }; + if root_events.slots.is_empty() || root_events.advances.is_empty() { + return Err(invalid( + "each pipeline state needs state_slot and state_advance operations", + )); + } + let roles = roles_by_make.entry(make).or_default(); + match state.role { + CutePipelineRoleAttr::Producer => { + roles.0 += 1; + if root_events.acquires.len() != 1 + || root_events.expects.len() != 1 + || root_events.advances.len() != 1 + || root_events.slots.len() != 1 + || root_events.tails.len() != 1 + || !root_events.waits.is_empty() + || !root_events.releases.is_empty() + { + return Err(invalid( + "producer pipeline state needs acquire, expect_tx, advance, and exactly one tail, with no consumer events", + )); + } + let acquire = root_events.acquires[0]; + let expect = root_events.expects[0]; + let advance = root_events.advances[0]; + let tail = root_events.tails[0]; + require_producer_pipeline_order(ctx, *root, acquire, expect, advance, tail)?; + let barrier = CutePipelineProducerExpectTxOp::wrap(expect).completion_barrier(ctx); + for r#use in barrier.uses(ctx) { + let copy = r#use.user_op(); + if !operation_dominates(ctx, expect, copy) + || !operation_dominates(ctx, copy, advance) + { + return Err(invalid( + "every TMA copy must be CFG-ordered after expect_tx and before producer state advance", + )); + } + } + } + CutePipelineRoleAttr::Consumer => { + roles.1 += 1; + if root_events.waits.len() != 1 + || root_events.releases.len() != 1 + || root_events.advances.len() != 1 + || root_events.slots.len() != 1 + || !root_events.acquires.is_empty() + || !root_events.expects.is_empty() + || !root_events.tails.is_empty() + { + return Err(invalid( + "consumer pipeline state needs wait, release, and advance, with no producer events", + )); + } + let wait = root_events.waits[0]; + let release = root_events.releases[0]; + let advance = root_events.advances[0]; + if !operation_dominates(ctx, *root, wait) + || !operation_dominates(ctx, wait, release) + || !operation_dominates(ctx, release, advance) + || !operation_reaches(ctx, advance, wait) + { + return Err(invalid( + "consumer pipeline CFG must order state_new -> wait -> release -> advance -> next wait", + )); + } + } + } + } + for make in makes { + if init_by_make + .get(&make) + .is_none_or(|initializers| initializers.len() != 1) + { + return Err(invalid( + "each tma_load_pipeline_make needs exactly one initialization", + )); + } + if roles_by_make.get(&make).copied() != Some((1, 1)) { + return Err(invalid( + "each tma_load_pipeline_make needs exactly one complete producer state and one complete consumer state", + )); + } + let initializer = init_by_make[&make][0]; + for (root, bound_make) in &binding { + if *bound_make == make && !operation_dominates(ctx, initializer, *root) { + return Err(invalid( + "load-pipeline initialization must dominate both producer and consumer state creation", + )); + } + } + } + Ok(()) +} diff --git a/cuda-oxide/cuteir/dialect-cute/src/verify/provenance.rs b/cuda-oxide/cuteir/dialect-cute/src/verify/provenance.rs new file mode 100644 index 0000000000..7bd4ae9a83 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/verify/provenance.rs @@ -0,0 +1,533 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Follow semantic origins through ordinary MIR carriers. +//! +//! Facts are sets of possible origins, including an explicit unknown origin. +//! Family analyses first propagate known facts to a fixed point, allowing loop +//! backedges to inherit their entry origin. They then seed unrelated values, +//! poison missing forwarding inputs, and propagate again before requiring one +//! unambiguous origin. Seeding unknowns earlier would reject valid loop state. +//! +//! A merge of an entry state and its own advanced state retains one origin; +//! merging either with an arbitrary scalar adds `Unknown` and is rejected. +//! Closed local cells conservatively union all stores, irrespective of order. +//! This tracks identity, not dynamic values or general memory aliasing. + +use std::collections::{HashMap, HashSet}; + +use dialect_mir::{ + attributes::MirCastKindAttr, + ops::{ + MirAllocaOp, MirCastOp, MirConstructStructOp, MirConstructTupleOp, MirExtractFieldOp, + MirFieldAddrOp, MirInsertFieldOp, MirLoadOp, MirStoreOp, + }, + types::{MirArrayType, MirPtrType, MirStructType, MirTupleType, address_space}, +}; +use pliron::{ + builtin::{op_interfaces::BranchOpInterface, types::IntegerType}, + context::{Context, Ptr}, + op::{Op, op_cast}, + operation::Operation, + r#type::{TypeHandle, Typed}, + value::Value, +}; + +/// Identity of a whole SSA carrier or one nested aggregate field. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub(super) struct CarrierKey { + /// The whole SSA carrier before following any aggregate field path. + pub(super) value: Value, + /// Empty means the value itself. Non-empty paths describe a field inside + /// an ordinary aggregate before a matching `mir.extract_field` exposes it. + pub(super) path: Vec, +} + +/// Monotone sets of possible semantic origins, indexed by carrier and field path. +#[derive(Clone)] +pub(super) struct OrdinaryFacts { + /// Accumulated origins; multiple entries at one key represent a merge that consumers must resolve. + pub(super) facts: HashMap>, +} + +impl Default for OrdinaryFacts { + fn default() -> Self { + Self { + facts: HashMap::new(), + } + } +} + +impl OrdinaryFacts +where + F: Clone + Eq + std::hash::Hash, +{ + /// Add one origin to a field path and report whether the fixed-point state changed. + pub(super) fn insert_path(&mut self, value: Value, path: Vec, fact: F) -> bool { + self.facts + .entry(CarrierKey { value, path }) + .or_default() + .insert(fact) + } + + /// Add an origin for the whole SSA carrier, with an empty field path. + pub(super) fn insert(&mut self, value: Value, fact: F) -> bool { + self.insert_path(value, Vec::new(), fact) + } + + /// Borrow the possible origins at an exact carrier field path. + pub(super) fn at_path(&self, value: Value, path: &[u32]) -> Option<&HashSet> { + self.facts.get(&CarrierKey { + value, + path: path.to_vec(), + }) + } + + /// Borrow origins attached directly to a whole SSA carrier. + pub(super) fn at(&self, value: Value) -> Option<&HashSet> { + self.at_path(value, &[]) + } + + /// Collect all nested field origins carried by one SSA value. + pub(super) fn entries(&self, value: Value) -> Vec<(Vec, F)> { + self.facts + .iter() + .filter(|(key, _)| key.value == value) + .flat_map(|(key, facts)| facts.iter().cloned().map(|fact| (key.path.clone(), fact))) + .collect() + } + + /// Forward origins below a source field prefix into the corresponding destination prefix. + pub(super) fn copy_entries( + &mut self, + source: Value, + destination: Value, + source_prefix: &[u32], + destination_prefix: &[u32], + ) -> bool { + let mut changed = false; + for (path, fact) in self.entries(source) { + let Some(suffix) = path.strip_prefix(source_prefix) else { + continue; + }; + let mut destination_path = destination_prefix.to_vec(); + destination_path.extend_from_slice(suffix); + changed |= self.insert_path(destination, destination_path, fact); + } + changed + } +} + +/// Read values passed to one successor through Pliron's branch interface. +pub(super) fn branch_successor_operands( + ctx: &Context, + terminator: Ptr, + successor_index: usize, +) -> Option> { + let dynamic = Operation::get_op_dyn(terminator, ctx); + let branch = op_cast::(dynamic.as_ref())?; + Some(branch.successor_operands(ctx, successor_index)) +} + +/// A provenance-preserving cast keeps the pointee and never strengthens +/// mutability. Local-cell aliases must also keep their address space; semantic +/// shared-memory carriers may use the compiler's exact generic/shared +/// round-trip, but no other address-space transition is accepted. +fn is_provenance_preserving_pointer_cast( + ctx: &Context, + operation: Ptr, + allow_shared_generic_transition: bool, +) -> bool { + if Operation::get_opid(operation, ctx) != MirCastOp::get_opid_static() + || operation.deref(ctx).get_num_operands() != 1 + || operation.deref(ctx).get_num_results() != 1 + || !MirCastOp::new(operation) + .get_attr_cast_kind(ctx) + .is_some_and(|kind| matches!(*kind, MirCastKindAttr::PtrToPtr)) + { + return false; + } + let operation_ref = operation.deref(ctx); + let source_type = operation_ref.get_operand(0).get_type(ctx); + let result_type = operation_ref.get_result(0).get_type(ctx); + let source_type_ref = source_type.deref(ctx); + let result_type_ref = result_type.deref(ctx); + let (Some(source), Some(result)) = ( + source_type_ref.downcast_ref::(), + result_type_ref.downcast_ref::(), + ) else { + return false; + }; + let address_space_is_preserved = source.address_space == result.address_space + || (allow_shared_generic_transition + && matches!( + (source.address_space, result.address_space), + (address_space::GENERIC, address_space::SHARED) + | (address_space::SHARED, address_space::GENERIC) + )); + source.pointee == result.pointee + && address_space_is_preserved + && (source.is_mutable || !result.is_mutable) +} + +/// Accept pointer casts preserving pointee, address space, and non-strengthened mutability. +pub(super) fn is_local_cell_pointer_alias_cast(ctx: &Context, operation: Ptr) -> bool { + is_provenance_preserving_pointer_cast(ctx, operation, false) +} + +/// Also permit exact generic/shared pointer round-trips for semantic shared carriers. +pub(super) fn is_semantic_pointer_carrier_cast(ctx: &Context, operation: Ptr) -> bool { + is_provenance_preserving_pointer_cast(ctx, operation, true) +} + +/// Resolve a pointer to one compiler-created local cell and an exact constant +/// field path. No pointer arithmetic, references, calls, or unknown casts are +/// accepted here. +pub(super) fn local_cell(ctx: &Context, mut pointer: Value) -> Option<(Value, Vec)> { + let mut reversed_path = Vec::new(); + let mut visited = HashSet::new(); + loop { + if !visited.insert(pointer) { + return None; + } + let definer = pointer.defining_op()?; + let opid = Operation::get_opid(definer, ctx); + let operation = definer.deref(ctx); + if opid == MirAllocaOp::get_opid_static() { + return (operation.get_num_results() == 1 && operation.get_result(0) == pointer).then( + || { + reversed_path.reverse(); + (pointer, reversed_path) + }, + ); + } + if opid == MirCastOp::get_opid_static() + && operation.get_num_operands() == 1 + && operation.get_num_results() == 1 + && operation.get_result(0) == pointer + && is_local_cell_pointer_alias_cast(ctx, definer) + { + pointer = operation.get_operand(0); + continue; + } + if opid == MirFieldAddrOp::get_opid_static() + && operation.get_num_operands() == 1 + && operation.get_num_results() == 1 + && operation.get_result(0) == pointer + { + let index = MirFieldAddrOp::new(definer).get_attr_field_index(ctx)?; + reversed_path.push(index.0); + pointer = operation.get_operand(0); + continue; + } + return None; + } +} + +/// Every alias of an accepted cell must stay within the same small +/// cast/field/load/store vocabulary. This lets the verifier conservatively +/// merge all stores to the cell without guessing about hidden aliases. +pub(super) fn local_cell_is_closed(ctx: &Context, root: Value) -> bool { + let mut worklist = vec![root]; + let mut visited = HashSet::new(); + while let Some(pointer) = worklist.pop() { + if !visited.insert(pointer) { + continue; + } + for r#use in pointer.uses(ctx) { + let user = r#use.user_op(); + let opid = Operation::get_opid(user, ctx); + let operation = user.deref(ctx); + if r#use.find_index(ctx) != 0 { + return false; + } + let forwards_pointer = operation.get_num_results() == 1 + && ((opid == MirCastOp::get_opid_static() + && is_local_cell_pointer_alias_cast(ctx, user)) + || (opid == MirFieldAddrOp::get_opid_static() + && MirFieldAddrOp::new(user) + .get_attr_field_index(ctx) + .is_some())); + if forwards_pointer { + worklist.push(operation.get_result(0)); + } else if opid == MirLoadOp::get_opid_static() { + if MirLoadOp::new(user).is_volatile(ctx) { + return false; + } + } else if opid == MirStoreOp::get_opid_static() { + if MirStoreOp::new(user).is_volatile(ctx) { + return false; + } + } else { + return false; + } + } + } + true +} + +/// Propagate semantic facts through the exact ordinary SSA adapters that can +/// remain after common CuTe preparation. This is deliberately a monotone +/// fixed-point transfer: loop block arguments become known once their entry +/// fact reaches them, then semantic advance operations can feed the same fact +/// around the backedge on a later iteration. +pub(super) fn propagate_ordinary_plumbing( + ctx: &Context, + operations: &[Ptr], + facts: &mut OrdinaryFacts, +) -> bool +where + F: Clone + Eq + std::hash::Hash, +{ + let mut changed = false; + let mut blocks = HashSet::new(); + for operation in operations { + let operation_ref = operation.deref(ctx); + if let Some(block) = operation_ref.get_parent_block() { + blocks.insert(block); + } + let opid = Operation::get_opid(*operation, ctx); + if opid == MirCastOp::get_opid_static() + && is_semantic_pointer_carrier_cast(ctx, *operation) + && operation_ref.get_num_operands() == 1 + && operation_ref.get_num_results() == 1 + { + changed |= facts.copy_entries( + operation_ref.get_operand(0), + operation_ref.get_result(0), + &[], + &[], + ); + } else if (opid == MirConstructStructOp::get_opid_static() + || opid == MirConstructTupleOp::get_opid_static()) + && operation_ref.get_num_results() == 1 + { + let result = operation_ref.get_result(0); + for (index, operand) in operation_ref.operands().enumerate() { + changed |= facts.copy_entries(operand, result, &[], &[index as u32]); + } + } else if opid == MirInsertFieldOp::get_opid_static() + && operation_ref.get_num_operands() == 2 + && operation_ref.get_num_results() == 1 + { + let Some(index) = MirInsertFieldOp::new(*operation).get_attr_insert_index(ctx) else { + continue; + }; + let base = operation_ref.get_operand(0); + let inserted = operation_ref.get_operand(1); + let result = operation_ref.get_result(0); + for (path, fact) in facts.entries(base) { + if path.first().copied() != Some(index.0) { + changed |= facts.insert_path(result, path, fact); + } + } + changed |= facts.copy_entries(inserted, result, &[], &[index.0]); + } else if opid == MirExtractFieldOp::get_opid_static() + && operation_ref.get_num_operands() == 1 + && operation_ref.get_num_results() == 1 + { + let Some(index) = MirExtractFieldOp::new(*operation).get_attr_index(ctx) else { + continue; + }; + changed |= facts.copy_entries( + operation_ref.get_operand(0), + operation_ref.get_result(0), + &[index.0], + &[], + ); + } else if opid == MirStoreOp::get_opid_static() && operation_ref.get_num_operands() == 2 { + let store = MirStoreOp::new(*operation); + if !store.is_volatile(ctx) + && let Some((root, path)) = local_cell(ctx, store.address_opd(ctx)) + && local_cell_is_closed(ctx, root) + { + changed |= facts.copy_entries(store.value_opd(ctx), root, &[], &path); + } + } else if opid == MirLoadOp::get_opid_static() + && operation_ref.get_num_operands() == 1 + && operation_ref.get_num_results() == 1 + { + let load = MirLoadOp::new(*operation); + if !load.is_volatile(ctx) + && let Some((root, path)) = local_cell(ctx, load.address_opd(ctx)) + && local_cell_is_closed(ctx, root) + { + changed |= facts.copy_entries(root, operation_ref.get_result(0), &path, &[]); + } + } + } + + for block in blocks { + let block_ref = block.deref(ctx); + for argument_index in 0..block_ref.get_num_arguments() { + let argument = block_ref.get_argument(argument_index); + for edge in block.uses(ctx) { + let terminator = edge.user_op(); + let Some(operands) = + branch_successor_operands(ctx, terminator, edge.find_index(ctx)) + else { + continue; + }; + if let Some(incoming) = operands.get(argument_index).copied() { + changed |= facts.copy_entries(incoming, argument, &[], &[]); + } + } + } + } + changed +} + +/// Enumerate integer fields of the requested bit width inside tuples, structs, and arrays. +fn integer_leaf_paths( + ctx: &Context, + ty: TypeHandle, + width: u32, + prefix: &mut Vec, + output: &mut Vec>, +) { + let ty_ref = ty.deref(ctx); + if ty_ref + .downcast_ref::() + .is_some_and(|integer| integer.width() == width) + { + output.push(prefix.clone()); + } else if let Some(structure) = ty_ref.downcast_ref::() { + for (index, field) in structure.field_types.iter().copied().enumerate() { + let Ok(index) = u32::try_from(index) else { + continue; + }; + prefix.push(index); + integer_leaf_paths(ctx, field, width, prefix, output); + prefix.pop(); + } + } else if let Some(tuple) = ty_ref.downcast_ref::() { + for (index, field) in tuple.get_types().iter().copied().enumerate() { + let Ok(index) = u32::try_from(index) else { + continue; + }; + prefix.push(index); + integer_leaf_paths(ctx, field, width, prefix, output); + prefix.pop(); + } + } else if let Some(array) = ty_ref.downcast_ref::() { + for index in 0..array.size() { + let Ok(index) = u32::try_from(index) else { + continue; + }; + prefix.push(index); + integer_leaf_paths(ctx, array.element_type(), width, prefix, output); + prefix.pop(); + } + } +} + +/// Seed factless integer leaves only after known semantic facts stabilize. +/// Aggregates are visited by field path, so an unrelated scalar cannot hide +/// inside a carrier that also contains known fields. +pub(super) fn seed_unknown_integer_carriers( + ctx: &Context, + operations: &[Ptr], + facts: &mut OrdinaryFacts, + width: u32, + unknown: impl Fn(Value) -> F, +) where + F: Clone + Eq + std::hash::Hash, +{ + let mut values = Vec::new(); + let mut blocks = HashSet::new(); + for operation in operations { + let operation_ref = operation.deref(ctx); + values.extend(operation_ref.results()); + if let Some(block) = operation_ref.get_parent_block() + && blocks.insert(block) + { + let block_ref = block.deref(ctx); + values.extend( + (0..block_ref.get_num_arguments()).map(|index| block_ref.get_argument(index)), + ); + } + } + + for value in values { + let mut paths = Vec::new(); + integer_leaf_paths(ctx, value.get_type(ctx), width, &mut Vec::new(), &mut paths); + for path in paths { + if facts.at_path(value, &path).is_none() { + facts.insert_path(value, path, unknown(value)); + } + } + } +} + +/// Mark missing inputs to a known block argument or closed local cell unknown. +/// Repeat this and ordinary propagation to a fixed point: one known incoming +/// edge must not conceal another edge or store with no semantic origin. +pub(super) fn poison_missing_forwarding_sources( + ctx: &Context, + operations: &[Ptr], + facts: &mut OrdinaryFacts, + unknown: impl Fn(Value) -> F + Copy, +) -> bool +where + F: Clone + Eq + std::hash::Hash, +{ + let mut changed = false; + let mut blocks = HashSet::new(); + for operation in operations { + if let Some(block) = operation.deref(ctx).get_parent_block() { + blocks.insert(block); + } + if Operation::get_opid(*operation, ctx) != MirStoreOp::get_opid_static() { + continue; + } + let store = MirStoreOp::new(*operation); + let Some((root, cell_path)) = local_cell(ctx, store.address_opd(ctx)) else { + continue; + }; + if !local_cell_is_closed(ctx, root) { + continue; + } + for (path, _) in facts.entries(root) { + let Some(source_path) = path.strip_prefix(cell_path.as_slice()) else { + continue; + }; + if facts.at_path(store.value_opd(ctx), source_path).is_none() { + changed |= facts.insert_path( + store.value_opd(ctx), + source_path.to_vec(), + unknown(store.value_opd(ctx)), + ); + } + } + } + + for block in blocks { + let block_ref = block.deref(ctx); + for argument_index in 0..block_ref.get_num_arguments() { + let argument = block_ref.get_argument(argument_index); + let argument_entries = facts.entries(argument); + if argument_entries.is_empty() { + continue; + } + for edge in block.uses(ctx) { + let terminator = edge.user_op(); + let Some(operands) = + branch_successor_operands(ctx, terminator, edge.find_index(ctx)) + else { + continue; + }; + let Some(incoming) = operands.get(argument_index).copied() else { + continue; + }; + for (path, _) in &argument_entries { + if facts.at_path(incoming, path).is_none() { + changed |= facts.insert_path(incoming, path.clone(), unknown(incoming)); + } + } + } + } + } + changed +} diff --git a/cuda-oxide/cuteir/dialect-cute/src/verify/scheduler.rs b/cuda-oxide/cuteir/dialect-cute/src/verify/scheduler.rs new file mode 100644 index 0000000000..2ca966e811 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/verify/scheduler.rs @@ -0,0 +1,382 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Persistent output-tile grids and their ordinary cursor provenance. + +use std::collections::{HashMap, HashSet}; + +use dialect_mir::ops::{MirNotOp, control_flow::MirCondBranchOp}; +use pliron::{ + basic_block::BasicBlock, + context::{Context, Ptr}, + op::{Op, OpId}, + operation::Operation, + r#type::Typed, + value::Value, +}; + +use super::{ + Result, + cfg::{block_reaches, operation_reaches}, + ghosts::scheduler_semantic_ids, + invalid, + provenance::{ + OrdinaryFacts, poison_missing_forwarding_sources, propagate_ordinary_plumbing, + seed_unknown_integer_carriers, + }, +}; +use crate::{ + attributes::CuteTileGridAttr, + scheduler_ops::{ + CuteSchedulerAdvanceOp, CuteSchedulerCurrentOp, CuteSchedulerHasWorkOp, + CuteSchedulerNew1dOp, CuteWorkTileCoordinatesOp, + }, + types::CuteWorkTileType, +}; + +/// Read the static output grid from a recognized scheduler operation. +fn scheduler_grid( + ctx: &Context, + operation: Ptr, + opid: &OpId, +) -> Result { + let grid = if *opid == CuteSchedulerNew1dOp::get_opid_static() { + CuteSchedulerNew1dOp::wrap(operation).tile_grid(ctx) + } else if *opid == CuteSchedulerHasWorkOp::get_opid_static() { + CuteSchedulerHasWorkOp::wrap(operation).tile_grid(ctx) + } else if *opid == CuteSchedulerCurrentOp::get_opid_static() { + CuteSchedulerCurrentOp::wrap(operation).tile_grid(ctx) + } else if *opid == CuteSchedulerAdvanceOp::get_opid_static() { + CuteSchedulerAdvanceOp::wrap(operation).tile_grid(ctx) + } else if *opid == CuteWorkTileCoordinatesOp::get_opid_static() { + CuteWorkTileCoordinatesOp::wrap(operation) + .work_tile(ctx) + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .map(|tile| tile.grid) + } else { + None + }; + grid.ok_or_else(|| invalid(format!("`{opid}` is missing its scheduler tile grid"))) +} + +/// Origin of an ordinary persistent-scheduler state component. +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +pub(super) enum SchedulerFact { + /// Linear tile cursor derived from one scheduler constructor. + Current(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Physical-grid stride derived from the same scheduler constructor. + Stride(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Scalar without scheduler provenance, invalid in a scheduler-state merge. + Unknown(#[doc = "Original SSA value lacking semantic provenance."] Value), +} + +/// Require exactly one current-index or stride origin of the requested kind. +pub(super) fn scheduler_root( + facts: &OrdinaryFacts, + value: Value, + current: bool, + what: &str, +) -> Result> { + let Some(found) = facts.at(value) else { + return Err(invalid(format!( + "{what} has no scheduler provenance after ordinary SSA forwarding" + ))); + }; + if found.len() != 1 { + return Err(invalid(format!( + "{what} merges values from different scheduler stories" + ))); + } + match found.iter().next().copied() { + Some(SchedulerFact::Current(root)) if current => Ok(root), + Some(SchedulerFact::Stride(root)) if !current => Ok(root), + _ => Err(invalid(format!( + "{what} is not derived from the required scheduler_new_1d result" + ))), + } +} + +/// Propagate scheduler origins through advances and ordinary carriers before poisoning unknown edges. +pub(super) fn build_scheduler_facts( + ctx: &Context, + operations: &[Ptr], +) -> OrdinaryFacts { + let mut facts = OrdinaryFacts::default(); + for operation in operations { + if Operation::get_opid(*operation, ctx) == CuteSchedulerNew1dOp::get_opid_static() { + let scheduler = CuteSchedulerNew1dOp::wrap(*operation); + facts.insert(scheduler.current(ctx), SchedulerFact::Current(*operation)); + facts.insert(scheduler.stride(ctx), SchedulerFact::Stride(*operation)); + } + } + + loop { + let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); + for operation in operations { + if Operation::get_opid(*operation, ctx) != CuteSchedulerAdvanceOp::get_opid_static() { + continue; + } + let advance = CuteSchedulerAdvanceOp::wrap(*operation); + let Some(current_facts) = facts.at(advance.current(ctx)).cloned() else { + continue; + }; + let Some(stride_facts) = facts.at(advance.stride(ctx)).cloned() else { + continue; + }; + for current in ¤t_facts { + let SchedulerFact::Current(root) = current else { + continue; + }; + if stride_facts.contains(&SchedulerFact::Stride(*root)) { + changed |= facts.insert(advance.next(ctx), SchedulerFact::Current(*root)); + } + } + } + if !changed { + break; + } + } + + // A factless scalar is an ordinary, unrelated value. Seeding it only + // after the semantic fixed point lets a valid loop recurrence establish + // itself first, while still poisoning a merge with an arbitrary edge. + seed_unknown_integer_carriers(ctx, operations, &mut facts, 64, SchedulerFact::Unknown); + loop { + let poisoned = + poison_missing_forwarding_sources(ctx, operations, &mut facts, SchedulerFact::Unknown); + let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); + if !poisoned && !propagated { + break; + } + } + facts +} + +/// Find a branch whose has-work condition guards the selected scheduler cursor. +fn scheduler_guard( + ctx: &Context, + has_work: Value, +) -> Option<(Ptr, Ptr, Ptr)> { + let mut value = has_work; + let mut negated = false; + loop { + let uses = value.uses(ctx); + if uses.len() != 1 || uses[0].find_index(ctx) != 0 { + return None; + } + let user = uses[0].user_op(); + let opid = Operation::get_opid(user, ctx); + if opid == MirNotOp::get_opid_static() { + let operation = user.deref(ctx); + if operation.get_num_operands() != 1 || operation.get_num_results() != 1 { + return None; + } + value = operation.get_result(0); + negated = !negated; + continue; + } + if opid != MirCondBranchOp::get_opid_static() { + return None; + } + let operation = user.deref(ctx); + if operation.get_num_successors() != 2 || operation.get_operand(0) != value { + return None; + } + let work_index = usize::from(negated); + return Some(( + user, + operation.get_successor(work_index), + operation.get_successor(1 - work_index), + )); + } +} + +/// Check one scheduler's state origins, loop advance, work guard, and coordinate consumers. +pub(super) fn verify_scheduler_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = scheduler_semantic_ids(); + let facts = build_scheduler_facts(ctx, all_ops); + let mut roots = HashSet::new(); + let mut has_work_roots = HashSet::new(); + let mut current_roots = HashSet::new(); + let mut advance_roots = HashSet::new(); + let mut coordinate_roots = HashSet::new(); + let mut has_work_by_root: HashMap, Vec>> = HashMap::new(); + let mut current_by_root: HashMap, Vec>> = HashMap::new(); + let mut advance_by_root: HashMap, Vec>> = HashMap::new(); + let mut coordinates_by_root: HashMap, Vec>> = HashMap::new(); + let mut currents = Vec::new(); + let mut coordinates = Vec::new(); + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if !ids.contains(&opid) { + continue; + } + if operation.deref(ctx).get_parent_block().is_none() { + return Err(invalid(format!( + "scheduler operation `{opid}` is not attached to a block" + ))); + } + let grid = scheduler_grid(ctx, *operation, &opid)?; + if opid == CuteSchedulerNew1dOp::get_opid_static() { + roots.insert(*operation); + } else if opid == CuteSchedulerHasWorkOp::get_opid_static() { + let semantic = CuteSchedulerHasWorkOp::wrap(*operation); + let root = scheduler_root( + &facts, + semantic.current(ctx), + true, + "cute.scheduler_has_work current", + )?; + if CuteSchedulerNew1dOp::wrap(root).tile_grid(ctx) != Some(grid) { + return Err(invalid( + "cute.scheduler_has_work grid differs from its scheduler_new_1d root", + )); + } + if !semantic.has_work(ctx).is_used(ctx) { + return Err(invalid( + "cute.scheduler_has_work result is unused; it must guard the scheduler body", + )); + } + has_work_roots.insert(root); + has_work_by_root.entry(root).or_default().push(*operation); + } else if opid == CuteSchedulerCurrentOp::get_opid_static() { + let semantic = CuteSchedulerCurrentOp::wrap(*operation); + let root = scheduler_root( + &facts, + semantic.current(ctx), + true, + "cute.scheduler_current current", + )?; + if CuteSchedulerNew1dOp::wrap(root).tile_grid(ctx) != Some(grid) { + return Err(invalid( + "cute.scheduler_current grid differs from its scheduler_new_1d root", + )); + } + current_roots.insert(root); + current_by_root.entry(root).or_default().push(*operation); + currents.push(*operation); + } else if opid == CuteWorkTileCoordinatesOp::get_opid_static() { + coordinates.push(*operation); + } else if opid == CuteSchedulerAdvanceOp::get_opid_static() { + let semantic = CuteSchedulerAdvanceOp::wrap(*operation); + let current_root = scheduler_root( + &facts, + semantic.current(ctx), + true, + "cute.scheduler_advance current", + )?; + let stride_root = scheduler_root( + &facts, + semantic.stride(ctx), + false, + "cute.scheduler_advance stride", + )?; + if current_root != stride_root + || CuteSchedulerNew1dOp::wrap(current_root).tile_grid(ctx) != Some(grid) + { + return Err(invalid( + "cute.scheduler_advance must use current, stride, and grid from one scheduler_new_1d", + )); + } + advance_roots.insert(current_root); + advance_by_root + .entry(current_root) + .or_default() + .push(*operation); + } + } + + for current in ¤ts { + let tile = CuteSchedulerCurrentOp::wrap(*current).work_tile(ctx); + let uses = tile.uses(ctx); + if uses.is_empty() + || uses.iter().any(|r#use| { + Operation::get_opid(r#use.user_op(), ctx) + != CuteWorkTileCoordinatesOp::get_opid_static() + || r#use.find_index(ctx) != 0 + }) + { + return Err(invalid( + "cute.scheduler_current work tile must be used only by at least one cute.work_tile_coordinates", + )); + } + } + for coordinates in coordinates { + let tile = CuteWorkTileCoordinatesOp::wrap(coordinates).work_tile(ctx); + let Some(producer) = tile.defining_op() else { + return Err(invalid( + "cute.work_tile_coordinates reached a block argument; a direct scheduler_current result is required", + )); + }; + if Operation::get_opid(producer, ctx) != CuteSchedulerCurrentOp::get_opid_static() + || !currents.contains(&producer) + { + return Err(invalid( + "cute.work_tile_coordinates must use a scheduler_current attached to this module", + )); + } + let root = scheduler_root( + &facts, + CuteSchedulerCurrentOp::wrap(producer).current(ctx), + true, + "cute.work_tile_coordinates scheduler current", + )?; + coordinate_roots.insert(root); + coordinates_by_root + .entry(root) + .or_default() + .push(coordinates); + } + + for root in roots { + if !has_work_roots.contains(&root) + || !current_roots.contains(&root) + || !coordinate_roots.contains(&root) + || !advance_roots.contains(&root) + { + return Err(invalid( + "each cute.scheduler_new_1d needs a complete has_work -> current -> coordinates -> advance story", + )); + } + let has_work = &has_work_by_root[&root]; + let current = ¤t_by_root[&root]; + let coordinates = &coordinates_by_root[&root]; + let advances = &advance_by_root[&root]; + if has_work.len() != 1 + || current.len() != 1 + || coordinates.len() != 1 + || advances.len() != 1 + { + return Err(invalid( + "each scheduler root needs exactly one static has_work/current/coordinates/advance protocol", + )); + } + let Some((guard, work_successor, exit_successor)) = + scheduler_guard(ctx, CuteSchedulerHasWorkOp::wrap(has_work[0]).has_work(ctx)) + else { + return Err(invalid( + "cute.scheduler_has_work must exclusively guard one conditional branch (optionally through mir.not)", + )); + }; + let current_block = current[0] + .deref(ctx) + .get_parent_block() + .expect("attached scheduler current"); + if !operation_reaches(ctx, root, has_work[0]) + || !block_reaches(ctx, work_successor, current_block) + || block_reaches(ctx, exit_successor, current_block) + || !operation_reaches(ctx, current[0], coordinates[0]) + || !operation_reaches(ctx, coordinates[0], advances[0]) + || !operation_reaches(ctx, advances[0], guard) + { + return Err(invalid( + "scheduler CFG must be new -> guarded work current -> coordinates -> advance -> next guard, with the exit edge outside the work body", + )); + } + } + Ok(()) +} diff --git a/cuda-oxide/cuteir/dialect-cute/src/verify/tensor.rs b/cuda-oxide/cuteir/dialect-cute/src/verify/tensor.rs new file mode 100644 index 0000000000..41d62b2e6a --- /dev/null +++ b/cuda-oxide/cuteir/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::{ + tensor_ops::{ + 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/cuteir/dialect-cute/src/verify/tests.rs b/cuda-oxide/cuteir/dialect-cute/src/verify/tests.rs new file mode 100644 index 0000000000..7496262072 --- /dev/null +++ b/cuda-oxide/cuteir/dialect-cute/src/verify/tests.rs @@ -0,0 +1,754 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression tests for CuTe operation, type, and semantic verification contracts. + +use dialect_mir::{ + attributes::MirCastKindAttr, + ops::{ + MirAllocaOp, MirCastOp, MirGotoOp, MirLoadOp, MirNotOp, MirReturnOp, MirStoreOp, + MirUndefOp, control_flow::MirCondBranchOp, function::MirFuncOp, + }, + types::{MirPtrType, address_space}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::{OperandSegmentInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + linked_list::ContainsLinkedList, + op::Op, + operation::Operation, + r#type::{TypeHandle, Typed}, + value::Value, +}; + +use super::{ + epilogue::require_epilogue_cfg_links, + mma::{build_mma_facts, require_smem_overlay_pair, verify_mma_use_closure}, + pipeline::require_producer_pipeline_order, + provenance::{ + is_local_cell_pointer_alias_cast, is_semantic_pointer_carrier_cast, local_cell_is_closed, + }, + scheduler::{build_scheduler_facts, scheduler_root}, + *, +}; +use crate::{ + attributes::{ + CutePipelineStateAttr, CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, + CuteTensorLayoutAttr, CuteTileGridAttr, CuteTiledMmaPlanAttr, + }, + gemm_tma_ops::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}, + layout::{ComposedLayout, OffsetUnit}, + pipeline_ops::{ + CutePipelineConsumerWaitOp, CutePipelineProducerExpectTxOp, CutePipelineStateNewOp, + CuteTmaLoadPipelineMakeOp, + }, + scheduler_ops::{ + CuteSchedulerAdvanceOp, CuteSchedulerCurrentOp, CuteSchedulerHasWorkOp, + CuteSchedulerNew1dOp, CuteWorkTileCoordinatesOp, + }, + smem_mma_ops::{CuteMmaPartitionBOp, CuteSmemTensorOverlayOp}, + types::CuteTensorViewType, +}; + +/// Create a test module and return its entry block for constructing semantic IR directly. +fn module_top(ctx: &mut Context) -> (Ptr, Ptr) { + dialect_mir::register(ctx); + crate::register(ctx); + let module = ModuleOp::new(ctx, "cute_verify_test".try_into().unwrap()); + let module = module.get_operation(); + let region = module.deref(ctx).get_region(0); + let block = region.deref(ctx).iter(ctx).next().unwrap(); + (module, block) +} + +/// Create a typed placeholder SSA result so a verifier test needs no runtime input. +fn undef(ctx: &mut Context, block: Ptr, ty: TypeHandle) -> Value { + let operation = MirUndefOp::new(ctx, ty).get_operation(); + operation.insert_at_back(block, ctx); + operation.deref(ctx).get_result(0) +} + +/// Construct a valid compiler-only contiguous tensor type for escape-audit tests. +fn ghost_tensor_type(ctx: &mut Context) -> TypeHandle { + let element: TypeHandle = FP32Type::get(ctx).into(); + CuteTensorViewType::get( + ctx, + element, + element, + CuteTensorAddressSpaceAttr::Gmem, + CuteTensorAccessAttr::ReadOnly, + 4, + CuteTensorLayoutAttr::Contiguous1D, + ) + .into() +} + +/// Append a compiler-owned local allocation used to test closed carrier storage. +fn local_cell(ctx: &mut Context, block: Ptr, ty: TypeHandle) -> Value { + let pointer: TypeHandle = MirPtrType::get_generic(ctx, ty, true).into(); + let alloca = Operation::new( + ctx, + MirAllocaOp::get_concrete_op_info(), + vec![pointer], + vec![], + vec![], + 0, + ); + alloca.insert_at_back(block, ctx); + alloca.deref(ctx).get_result(0) +} + +/// Append a non-volatile MIR store into the fixture's local cell. +fn store(ctx: &mut Context, block: Ptr, pointer: Value, value: Value) { + Operation::new( + ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![pointer, value], + vec![], + 0, + ) + .insert_at_back(block, ctx); +} + +/// Append a non-volatile MIR load and return its SSA result. +fn load(ctx: &mut Context, block: Ptr, pointer: Value, ty: TypeHandle) -> Value { + let load = Operation::new( + ctx, + MirLoadOp::get_concrete_op_info(), + vec![ty], + vec![pointer], + vec![], + 0, + ); + load.insert_at_back(block, ctx); + load.deref(ctx).get_result(0) +} + +/// Append a pointer cast with controlled pointee, mutability, and address-space facts. +fn pointer_cast( + ctx: &mut Context, + block: Ptr, + pointer: Value, + result_type: TypeHandle, +) -> Ptr { + let cast = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![result_type], + vec![pointer], + vec![], + 0, + ); + MirCastOp::new(cast).set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); + cast.insert_at_back(block, ctx); + cast +} + +/// Append an unconditional successor edge with explicit block-argument operands. +fn goto(ctx: &mut Context, block: Ptr, target: Ptr, operands: Vec) { + Operation::new( + ctx, + MirGotoOp::get_concrete_op_info(), + vec![], + operands, + vec![target], + 0, + ) + .insert_at_back(block, ctx); +} + +/// Append a conditional branch with distinct successor operands for CFG/provenance tests. +fn cond_branch( + ctx: &mut Context, + block: Ptr, + condition: Value, + true_target: Ptr, + false_target: Ptr, +) { + let (operands, sizes) = + MirCondBranchOp::compute_segment_sizes(vec![vec![condition], Vec::new(), Vec::new()]); + let branch = Operation::new( + ctx, + MirCondBranchOp::get_concrete_op_info(), + vec![], + operands, + vec![true_target, false_target], + 0, + ); + MirCondBranchOp::new(branch).set_operand_segment_sizes(ctx, sizes); + branch.insert_at_back(block, ctx); +} + +/// Construct one valid persistent-scheduler function with configurable body/exit placement. +fn append_complete_scheduler_function( + ctx: &mut Context, + module_block: Ptr, + name: &str, + grid: CuteTileGridAttr, + misroute_work_edge: bool, +) { + let function_type = FunctionType::get(ctx, vec![], vec![]); + let function = Operation::new( + ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function_op = MirFuncOp::new(ctx, function, TypeAttr::new(function_type.into())); + function_op.set_symbol_name(ctx, name.try_into().unwrap()); + function.insert_at_back(module_block, ctx); + let region = function.deref(ctx).get_region(0); + let u64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + let entry = BasicBlock::new(ctx, None, vec![]); + entry.insert_at_back(region, ctx); + let header = BasicBlock::new(ctx, None, vec![u64_ty, u64_ty]); + header.insert_at_back(region, ctx); + let body = BasicBlock::new(ctx, None, vec![]); + body.insert_at_back(region, ctx); + let exit = BasicBlock::new(ctx, None, vec![]); + exit.insert_at_back(region, ctx); + + let start = CuteSchedulerNew1dOp::new(ctx, grid).get_operation(); + start.insert_at_back(entry, ctx); + let start = CuteSchedulerNew1dOp::wrap(start); + goto( + ctx, + entry, + header, + vec![start.current(ctx), start.stride(ctx)], + ); + + let current = header.deref(ctx).get_argument(0); + let stride = header.deref(ctx).get_argument(1); + let has_work = CuteSchedulerHasWorkOp::new(ctx, current, grid).get_operation(); + has_work.insert_at_back(header, ctx); + let not = Operation::new( + ctx, + MirNotOp::get_concrete_op_info(), + vec![ + CuteSchedulerHasWorkOp::wrap(has_work) + .has_work(ctx) + .get_type(ctx), + ], + vec![CuteSchedulerHasWorkOp::wrap(has_work).has_work(ctx)], + vec![], + 0, + ); + not.insert_at_back(header, ctx); + let finished = not.deref(ctx).get_result(0); + let (true_target, false_target) = if misroute_work_edge { + (body, exit) + } else { + (exit, body) + }; + cond_branch(ctx, header, finished, true_target, false_target); + + let selected = CuteSchedulerCurrentOp::new(ctx, current, grid).get_operation(); + selected.insert_at_back(body, ctx); + CuteWorkTileCoordinatesOp::new(ctx, CuteSchedulerCurrentOp::wrap(selected).work_tile(ctx)) + .get_operation() + .insert_at_back(body, ctx); + let advance = CuteSchedulerAdvanceOp::new(ctx, current, stride, grid).get_operation(); + advance.insert_at_back(body, ctx); + goto( + ctx, + body, + header, + vec![CuteSchedulerAdvanceOp::wrap(advance).next(ctx), stride], + ); + + Operation::new( + ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ) + .insert_at_back(exit, ctx); +} + +/// The escape audit must inspect ordinary IR types even when no semantic CuTe operation exposes the ghost. +#[test] +fn rejects_a_ghost_nested_in_an_ordinary_pointer_without_any_cute_op() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + let ghost = ghost_tensor_type(&mut ctx); + let inner: TypeHandle = MirPtrType::get_generic(&mut ctx, ghost, false).into(); + let outer: TypeHandle = MirPtrType::get_generic(&mut ctx, inner, false).into(); + undef(&mut ctx, block, outer); + + let error = verify_cute_semantics(&ctx, module).unwrap_err().to_string(); + assert!(error.contains("ghost CuTe value"), "{error}"); + assert!(error.contains("operand or result"), "{error}"); +} + +/// Function scopes isolate scheduler origins, and verification must leave the input module unchanged. +#[test] +fn accepts_independent_scheduler_stories_in_two_functions_without_mutation() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + append_complete_scheduler_function( + &mut ctx, + block, + "first", + CuteTileGridAttr::new(4, 4, 1), + false, + ); + append_complete_scheduler_function( + &mut ctx, + block, + "second", + CuteTileGridAttr::new(5, 4, 1), + false, + ); + let mut before = Vec::new(); + collect_ops(&ctx, module, &mut before); + + verify_cute_semantics(&ctx, module).unwrap(); + + let mut after = Vec::new(); + collect_ops(&ctx, module, &mut after); + assert_eq!(after, before); +} + +/// CFG work guards, rather than textual traversal order, determine whether a selected tile is valid. +#[test] +fn rejects_a_scheduler_body_on_the_exit_edge_even_when_preorder_looks_ordered() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + append_complete_scheduler_function( + &mut ctx, + block, + "misrouted", + CuteTileGridAttr::new(4, 4, 1), + true, + ); + + let error = verify_cute_semantics(&ctx, module).unwrap_err().to_string(); + assert!(error.contains("scheduler CFG"), "{error}"); +} + +/// Non-escaping compiler-owned loads/stores may forward a scheduler origin without requiring full SSA promotion. +#[test] +fn accepts_a_closed_local_cell_carrying_scheduler_provenance() { + let mut ctx = Context::new(); + module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let block = BasicBlock::new(&mut ctx, None, vec![]); + let grid = CuteTileGridAttr::new(4, 4, 1); + let start = CuteSchedulerNew1dOp::new(&mut ctx, grid).get_operation(); + start.insert_at_back(block, &ctx); + let pointer = local_cell(&mut ctx, block, u64_ty); + let current = CuteSchedulerNew1dOp::wrap(start).current(&ctx); + store(&mut ctx, block, pointer, current); + let carried = load(&mut ctx, block, pointer, u64_ty); + let operations: Vec<_> = block.deref(&ctx).iter(&ctx).collect(); + + let facts = build_scheduler_facts(&ctx, &operations); + assert_eq!( + scheduler_root(&facts, carried, true, "test carrier").unwrap(), + start + ); +} + +/// One unrelated incoming store must invalidate a cell even when another store has valid scheduler provenance. +#[test] +fn factless_block_argument_store_poisons_a_valid_scheduler_cell() { + let mut ctx = Context::new(); + module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let block = BasicBlock::new(&mut ctx, None, vec![u64_ty]); + let arbitrary = block.deref(&ctx).get_argument(0); + let grid = CuteTileGridAttr::new(4, 4, 1); + let start = CuteSchedulerNew1dOp::new(&mut ctx, grid).get_operation(); + start.insert_at_back(block, &ctx); + let pointer = local_cell(&mut ctx, block, u64_ty); + let current = CuteSchedulerNew1dOp::wrap(start).current(&ctx); + store(&mut ctx, block, pointer, current); + store(&mut ctx, block, pointer, arbitrary); + let poisoned = load(&mut ctx, block, pointer, u64_ty); + let operations: Vec<_> = block.deref(&ctx).iter(&ctx).collect(); + + let facts = build_scheduler_facts(&ctx, &operations); + let error = scheduler_root(&facts, poisoned, true, "test carrier") + .unwrap_err() + .to_string(); + assert!(error.contains("merges values"), "{error}"); +} + +/// Casts that change pointee, address space, or write capability cannot establish closed-cell aliases. +#[test] +fn non_identity_pointer_casts_cannot_alias_a_local_semantic_cell() { + let mut ctx = Context::new(); + module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let block = BasicBlock::new(&mut ctx, None, vec![]); + let pointer = local_cell(&mut ctx, block, u64_ty); + let shared_pointer_ty: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::SHARED).into(); + let different_pointee_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, true).into(); + for result_type in [shared_pointer_ty, different_pointee_ty] { + let cast = pointer_cast(&mut ctx, block, pointer, result_type); + let alias = cast.deref(&ctx).get_result(0); + assert!(!is_local_cell_pointer_alias_cast(&ctx, cast)); + assert!(super::provenance::local_cell(&ctx, alias).is_none()); + } + let read_only_type: TypeHandle = MirPtrType::get_generic(&mut ctx, u64_ty, false).into(); + let read_only = pointer_cast(&mut ctx, block, pointer, read_only_type); + let read_only_alias = read_only.deref(&ctx).get_result(0); + assert!(is_local_cell_pointer_alias_cast(&ctx, read_only)); + assert!(super::provenance::local_cell(&ctx, read_only_alias).is_some()); + + let mutable_type: TypeHandle = MirPtrType::get_generic(&mut ctx, u64_ty, true).into(); + let escalation = pointer_cast(&mut ctx, block, read_only_alias, mutable_type); + let escalated_alias = escalation.deref(&ctx).get_result(0); + assert!(!is_local_cell_pointer_alias_cast(&ctx, escalation)); + assert!(super::provenance::local_cell(&ctx, escalated_alias).is_none()); + assert!(!local_cell_is_closed(&ctx, pointer)); +} + +/// Semantic shared carriers may survive the exact generic/shared round-trip while other address-space casts are rejected. +#[test] +fn semantic_pointer_carriers_allow_only_shared_generic_address_space_roundtrips() { + let mut ctx = Context::new(); + module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let block = BasicBlock::new(&mut ctx, None, vec![]); + let shared_type: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::SHARED).into(); + let generic_type: TypeHandle = MirPtrType::get_generic(&mut ctx, u64_ty, true).into(); + let shared = undef(&mut ctx, block, shared_type); + let to_generic = pointer_cast(&mut ctx, block, shared, generic_type); + assert!(is_semantic_pointer_carrier_cast(&ctx, to_generic)); + let generic = to_generic.deref(&ctx).get_result(0); + let to_shared = pointer_cast(&mut ctx, block, generic, shared_type); + assert!(is_semantic_pointer_carrier_cast(&ctx, to_shared)); + + let global_type: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::GLOBAL).into(); + let global = undef(&mut ctx, block, global_type); + let global_to_generic = pointer_cast(&mut ctx, block, global, generic_type); + assert!(!is_semantic_pointer_carrier_cast(&ctx, global_to_generic)); + + let generic_u32_type: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, true).into(); + let wrong_pointee = pointer_cast(&mut ctx, block, shared, generic_u32_type); + assert!(!is_semantic_pointer_carrier_cast(&ctx, wrong_pointee)); + + let read_only_shared_type: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, false, address_space::SHARED).into(); + let read_only_shared = undef(&mut ctx, block, read_only_shared_type); + let mutability_escalation = pointer_cast(&mut ctx, block, read_only_shared, generic_type); + assert!(!is_semantic_pointer_carrier_cast( + &ctx, + mutability_escalation + )); +} + +/// Matching scalar types cannot substitute a consumer-origin cursor for producer lifecycle state. +#[test] +fn rejects_pipeline_lifecycle_state_from_the_wrong_role_root() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let barrier_ty: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::SHARED).into(); + let barrier_base = undef(&mut ctx, block, barrier_ty); + let make = CuteTmaLoadPipelineMakeOp::new(&mut ctx, barrier_base, 3, 8, 8, 8).get_operation(); + make.insert_at_back(block, &ctx); + let pipeline = CuteTmaLoadPipelineMakeOp::wrap(make).pipeline(&ctx); + let producer = + CutePipelineStateNewOp::new(&mut ctx, CutePipelineStateAttr::producer(3)).get_operation(); + producer.insert_at_back(block, &ctx); + let producer = CutePipelineStateNewOp::wrap(producer); + let producer_slot = producer.slot(&ctx); + let producer_phase = producer.phase(&ctx); + CutePipelineConsumerWaitOp::new( + &mut ctx, + pipeline, + producer_slot, + producer_phase, + CutePipelineStateAttr::consumer(3), + ) + .get_operation() + .insert_at_back(block, &ctx); + + let error = verify_cute_semantics(&ctx, module).unwrap_err().to_string(); + assert!(error.contains("state facts differ"), "{error}"); +} + +/// Producer drain must follow the cursor's acquire/expect/advance chain in the CFG. +#[test] +fn rejects_producer_tail_that_precedes_its_lifecycle() { + let mut ctx = Context::new(); + module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let block = BasicBlock::new(&mut ctx, None, vec![]); + let root = undef(&mut ctx, block, u64_ty).defining_op().unwrap(); + let tail = undef(&mut ctx, block, u64_ty).defining_op().unwrap(); + let acquire = undef(&mut ctx, block, u64_ty).defining_op().unwrap(); + let expect = undef(&mut ctx, block, u64_ty).defining_op().unwrap(); + let advance = undef(&mut ctx, block, u64_ty).defining_op().unwrap(); + + let error = require_producer_pipeline_order(&ctx, root, acquire, expect, advance, tail) + .unwrap_err() + .to_string(); + assert!(error.contains("producer pipeline CFG"), "{error}"); +} + +/// A correctly typed barrier still needs provenance from the matching transaction-expectation operation. +#[test] +fn rejects_tma_copy_with_an_arbitrary_typed_completion_barrier() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let barrier_ty: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::SHARED).into(); + let arbitrary_barrier = undef(&mut ctx, block, barrier_ty); + let layout = + ComposedLayout::from_layout("(1,16):(16,1)".parse().unwrap(), OffsetUnit::Elements); + let descriptor_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); + let shared_ty: TypeHandle = + MirPtrType::get(&mut ctx, u8_ty, true, address_space::SHARED).into(); + let descriptor = undef(&mut ctx, block, descriptor_ty); + let shared = undef(&mut ctx, block, shared_ty); + let capacity = undef(&mut ctx, block, u64_ty); + let row = undef(&mut ctx, block, u64_ty); + let column = undef(&mut ctx, block, u64_ty); + let source = + CuteTmaGmemViewOp::new(&mut ctx, descriptor, u8_ty, layout.clone()).get_operation(); + source.insert_at_back(block, &ctx); + let destination = + CuteTmaSmemViewOp::new(&mut ctx, shared, capacity, u8_ty, layout, 16).get_operation(); + destination.insert_at_back(block, &ctx); + let source_view = source.deref(&ctx).get_result(0); + let destination_view = destination.deref(&ctx).get_result(0); + CuteTmaCopy2dOp::new( + &mut ctx, + source_view, + destination_view, + row, + column, + arbitrary_barrier, + ) + .get_operation() + .insert_at_back(block, &ctx); + + let error = verify_cute_semantics(&ctx, module).unwrap_err().to_string(); + assert!( + error.contains("completion barrier is not the result of an expect_tx"), + "{error}" + ); +} + +/// Shared base and capacity must share one overlay origin, even when both overlays have identical types. +#[test] +fn rejects_mma_base_and_capacity_from_unrelated_same_typed_overlays() { + let mut ctx = Context::new(); + let (_module, block) = module_top(&mut ctx); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let shared_ty: TypeHandle = + MirPtrType::get(&mut ctx, u8_ty, true, address_space::SHARED).into(); + let first_base = undef(&mut ctx, block, shared_ty); + let first_capacity = undef(&mut ctx, block, u64_ty); + let second_base = undef(&mut ctx, block, shared_ty); + let second_capacity = undef(&mut ctx, block, u64_ty); + let first = + CuteSmemTensorOverlayOp::new(&mut ctx, first_base, first_capacity, u64_ty).get_operation(); + first.insert_at_back(block, &ctx); + let second = CuteSmemTensorOverlayOp::new(&mut ctx, second_base, second_capacity, u64_ty) + .get_operation(); + second.insert_at_back(block, &ctx); + let operations: Vec<_> = block.deref(&ctx).iter(&ctx).collect(); + let facts = build_mma_facts(&ctx, &operations); + + let error = require_smem_overlay_pair( + &facts, + CuteSmemTensorOverlayOp::wrap(first).base(&ctx), + CuteSmemTensorOverlayOp::wrap(second).capacity(&ctx), + "test MMA carrier", + ) + .unwrap_err() + .to_string(); + assert!(error.contains("different shared overlays"), "{error}"); +} + +/// One valid semantic use does not permit another use to escape the overlay's closed carrier flow. +#[test] +fn rejects_mma_overlay_with_a_valid_partition_edge_and_an_unrelated_load() { + let mut ctx = Context::new(); + let (_module, block) = module_top(&mut ctx); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let shared_ty: TypeHandle = + MirPtrType::get(&mut ctx, u8_ty, true, address_space::SHARED).into(); + let base = undef(&mut ctx, block, shared_ty); + let capacity = undef(&mut ctx, block, u64_ty); + let overlay = CuteSmemTensorOverlayOp::new(&mut ctx, base, capacity, u64_ty).get_operation(); + overlay.insert_at_back(block, &ctx); + let overlay = CuteSmemTensorOverlayOp::wrap(overlay); + let overlay_base = overlay.base(&ctx); + let overlay_capacity = overlay.capacity(&ctx); + let warp_n = undef(&mut ctx, block, u64_ty); + let k_half = undef(&mut ctx, block, u64_ty); + let placement = + ComposedLayout::from_layout("(128,32):(32,1)".parse().unwrap(), OffsetUnit::Elements); + CuteMmaPartitionBOp::new( + &mut ctx, + overlay_base, + overlay_capacity, + warp_n, + k_half, + u64_ty, + CuteTiledMmaPlanAttr::mxf4_128x128x128(placement), + ) + .get_operation() + .insert_at_back(block, &ctx); + Operation::new( + &mut ctx, + MirLoadOp::get_concrete_op_info(), + vec![u8_ty], + vec![overlay_base], + vec![], + 0, + ) + .insert_at_back(block, &ctx); + let operations: Vec<_> = block.deref(&ctx).iter(&ctx).collect(); + let facts = build_mma_facts(&ctx, &operations); + + let error = verify_mma_use_closure(&ctx, &operations, &facts) + .unwrap_err() + .to_string(); + assert!(error.contains("unsupported outgoing use"), "{error}"); + assert!(error.contains("mir.load"), "{error}"); +} + +/// A plausible traversal order cannot replace a reachable writer-to-issuer synchronization path. +#[test] +fn rejects_epilogue_writer_and_issuer_in_mutually_exclusive_preorder_blocks() { + let mut ctx = Context::new(); + let (_module, module_block) = module_top(&mut ctx); + let function_type = FunctionType::get(&ctx, vec![], vec![]); + let function = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(&mut ctx, function, TypeAttr::new(function_type.into())) + .set_symbol_name(&mut ctx, "epilogue_cfg".try_into().unwrap()); + function.insert_at_back(module_block, &ctx); + let region = function.deref(&ctx).get_region(0); + let entry = BasicBlock::new(&mut ctx, None, vec![]); + entry.insert_at_back(region, &ctx); + let writer = BasicBlock::new(&mut ctx, None, vec![]); + writer.insert_at_back(region, &ctx); + let issuer = BasicBlock::new(&mut ctx, None, vec![]); + issuer.insert_at_back(region, &ctx); + let i1_ty: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let condition = undef(&mut ctx, entry, i1_ty); + cond_branch(&mut ctx, entry, condition, writer, issuer); + let acquire = undef(&mut ctx, writer, u64_ty).defining_op().unwrap(); + let reusable = undef(&mut ctx, writer, u64_ty).defining_op().unwrap(); + let ready = undef(&mut ctx, writer, u64_ty).defining_op().unwrap(); + Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ) + .insert_at_back(writer, &ctx); + let half_zero = undef(&mut ctx, issuer, u64_ty).defining_op().unwrap(); + let commit = undef(&mut ctx, issuer, u64_ty).defining_op().unwrap(); + let tail = undef(&mut ctx, issuer, u64_ty).defining_op().unwrap(); + Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ) + .insert_at_back(issuer, &ctx); + + let error = require_epilogue_cfg_links(&ctx, acquire, reusable, ready, half_zero, commit, tail) + .unwrap_err() + .to_string(); + assert!(error.contains("epilogue CFG"), "{error}"); +} + +/// The expectation promise must equal the sum of its direct TMA users' static tile bytes. +#[test] +fn rejects_expect_tx_when_direct_tma_bytes_do_not_match() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let barrier_ty: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::SHARED).into(); + let barrier_base = undef(&mut ctx, block, barrier_ty); + let make = CuteTmaLoadPipelineMakeOp::new(&mut ctx, barrier_base, 3, 8, 8, 8).get_operation(); + make.insert_at_back(block, &ctx); + let pipeline = CuteTmaLoadPipelineMakeOp::wrap(make).pipeline(&ctx); + let state = CutePipelineStateAttr::producer(3); + let state_root = CutePipelineStateNewOp::new(&mut ctx, state).get_operation(); + state_root.insert_at_back(block, &ctx); + let slot = CutePipelineStateNewOp::wrap(state_root).slot(&ctx); + let expect = + CutePipelineProducerExpectTxOp::new(&mut ctx, pipeline, slot, state).get_operation(); + expect.insert_at_back(block, &ctx); + + let layout = + ComposedLayout::from_layout("(1,16):(16,1)".parse().unwrap(), OffsetUnit::Elements); + let descriptor_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); + let shared_ty: TypeHandle = + MirPtrType::get(&mut ctx, u8_ty, true, address_space::SHARED).into(); + let descriptor = undef(&mut ctx, block, descriptor_ty); + let shared = undef(&mut ctx, block, shared_ty); + let capacity = undef(&mut ctx, block, u64_ty); + let row = undef(&mut ctx, block, u64_ty); + let column = undef(&mut ctx, block, u64_ty); + let source = + CuteTmaGmemViewOp::new(&mut ctx, descriptor, u8_ty, layout.clone()).get_operation(); + source.insert_at_back(block, &ctx); + let destination = + CuteTmaSmemViewOp::new(&mut ctx, shared, capacity, u8_ty, layout, 16).get_operation(); + destination.insert_at_back(block, &ctx); + let source_view = source.deref(&ctx).get_result(0); + let destination_view = destination.deref(&ctx).get_result(0); + let completion = CutePipelineProducerExpectTxOp::wrap(expect).completion_barrier(&ctx); + CuteTmaCopy2dOp::new( + &mut ctx, + source_view, + destination_view, + row, + column, + completion, + ) + .get_operation() + .insert_at_back(block, &ctx); + + let error = verify_cute_semantics(&ctx, module).unwrap_err().to_string(); + assert!(error.contains("promises 8 bytes"), "{error}"); + assert!(error.contains("move 16 bytes"), "{error}"); +} diff --git a/cuda-oxide/cuteir/dialect-cute/src/verify/tma.rs b/cuda-oxide/cuteir/dialect-cute/src/verify/tma.rs new file mode 100644 index 0000000000..baba41dc70 --- /dev/null +++ b/cuda-oxide/cuteir/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::{ + epilogue_ops::CuteTmaStore2dSemanticOp, + gemm_tma_ops::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}, + pipeline_ops::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/cuteir/dialect-cute/tests/global_store.rs b/cuda-oxide/cuteir/dialect-cute/tests/global_store.rs index 24548e8e8d..952756ad1a 100644 --- a/cuda-oxide/cuteir/dialect-cute/tests/global_store.rs +++ b/cuda-oxide/cuteir/dialect-cute/tests/global_store.rs @@ -1,3 +1,4 @@ +//! 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 @@ -11,6 +12,7 @@ use pliron::{ 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(); diff --git a/cuda-oxide/cuteir/dialect-cute/tests/sm100_attention.rs b/cuda-oxide/cuteir/dialect-cute/tests/sm100_attention.rs index 66568fc848..78578eaad1 100644 --- a/cuda-oxide/cuteir/dialect-cute/tests/sm100_attention.rs +++ b/cuda-oxide/cuteir/dialect-cute/tests/sm100_attention.rs @@ -2,8 +2,12 @@ * 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::sm100_attention_ops::{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, @@ -16,6 +20,7 @@ fn qk() -> CuteSm100MmaAttr { 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(); @@ -32,6 +37,7 @@ fn independent_qk_pv_layouts_verify() { }; 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(); @@ -45,6 +51,7 @@ fn rejects_wrong_group_shape_and_major() { 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()); @@ -52,6 +59,7 @@ fn rejects_invalid_shared_layout_extents() { 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; @@ -87,6 +95,7 @@ fn tma_rejects_noncanonical_layout() { ); } +/// A global-to-shared transfer must retain writable destination capability. #[test] fn tma_load_rejects_readonly_shared_destination() { use dialect_cute::sm100_attention_ops::{CuteSm100TmaAttr, CuteSm100TmaLoadOp}; @@ -125,6 +134,7 @@ fn tma_load_rejects_readonly_shared_destination() { } } +/// 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::sm100_attention_ops::{CuteSm100MmaSsOp, CuteSm100MmaTsOp}; diff --git a/cuda-oxide/cuteir/docs/translation.md b/cuda-oxide/cuteir/docs/translation.md index a0093e0001..5d30c23a62 100644 --- a/cuda-oxide/cuteir/docs/translation.md +++ b/cuda-oxide/cuteir/docs/translation.md @@ -10,7 +10,7 @@ Compile Rust CuTe kernels with NVIDIA's pinned CUTLASS 4.7 compiler: ## Install and build -From the `cuda-oxide/` directory after normal CUDA Oxide setup: +From the `cuda-oxide/` directory after normal cuda-oxide setup: ```bash cargo oxide toolchain install cutlass @@ -55,7 +55,7 @@ hash; changing the environment option invalidates the device build fingerprint. | [Block-scaled GEMM](../examples/blockscale_gemm_cute) | Scheduler, work tiles, TMA pipelines, shared-memory MMA, epilogue stores | | [SM100 FP16 GEMM](../examples/fp16_gemm_256x352_cute) | Typed shared tiles, TMEM, paired two-CTA MMA, cluster TMA, producer/consumer pipelines | -The [SM100 attention API](../cute-rs/src/sm100_attention.rs) supplies independent +The [SM100 attention API](../cute-rs/src/sm100_attention/mod.rs) supplies independent QK and PV plans: group-one/group-two shared/shared MMA, group-one TMEM/shared MMA, K-major and MN-major tiled shared layouts, rank-three typed TMA copies, and FP32 TMEM fragments with packed FP16 stores. The semantic verifier checks @@ -84,7 +84,7 @@ for the full launch contract. ## SM100 contracts -The [Rust APIs](../cute-rs/src/sm100.rs) become verified `cute.sm100_*` plans: +The [Rust APIs](../cute-rs/src/sm100/mod.rs) become verified `cute.sm100_*` plans: | API | Contract | | --- | --- | diff --git a/cuda-oxide/cuteir/examples/blockscale_gemm_ctm/src/kernel.rs b/cuda-oxide/cuteir/examples/blockscale_gemm_ctm/src/kernel.rs index d19514c891..3cb0c4a823 100644 --- a/cuda-oxide/cuteir/examples/blockscale_gemm_ctm/src/kernel.rs +++ b/cuda-oxide/cuteir/examples/blockscale_gemm_ctm/src/kernel.rs @@ -13,28 +13,48 @@ use ctm_rs::{ tma::{self, Coordinates2D, TensorMap}, }; use cuda_device::{DynamicSharedArray, cuda_module, kernel, launch_bounds, launch_contract}; +/// Output row count for the fixed benchmark problem. pub const M: usize = 2048; +/// Output column count for the fixed benchmark problem. pub const N: usize = 2048; +/// Reduction extent in logical FP4 elements before packing. pub const K: usize = 1024; +/// Number of 32-element K groups, each with one UE8M0 scale. pub const KB: usize = K / 32; +/// Number of 128-element K tiles consumed for each output tile. pub const K_STAGES: usize = K / 128; +/// Threads required in each CTA by the kernel launch contract. pub const THREADS: usize = 288; +/// Number of 128x128 output tiles assigned persistently across the launched CTAs. pub const TOTAL_TILES: usize = (M / 128) * (N / 128); +/// Total dynamic shared-memory allocation in bytes required per CTA. pub const DYNAMIC_SMEM_BYTES: usize = 86_016; +/// Number of shared input slots overlapped by TMA and matrix instructions. const STAGES: usize = 3; +/// Byte footprint of one 128-row packed FP4 value stage. const DATA_BYTES: usize = 8192; +/// Byte footprint of one 128-row, four-group scale stage. const SCALE_BYTES: usize = 512; +/// Byte offset of the staged B values after all A stages. const B_OFFSET: usize = 24576; +/// Byte offset of the A-scale ring after both value rings. const SFA_OFFSET: usize = 49152; +/// Byte offset of the B-scale ring after the A-scale ring. const SFB_OFFSET: usize = 50688; +/// Byte offset of the full/empty barrier rings in dynamic shared memory. const BARRIERS: usize = 52224; +/// Aligned byte offset of the two shared FP16 output halves. const OUTPUT: usize = 53248; +/// Total value and scale bytes expected by one full-stage barrier. const TX_BYTES: u32 = 17408; +/// Compile-time check that the two output halves end at the declared shared-memory size. const _: () = assert!(OUTPUT + 2 * F16Tile128::FOOTPRINT_BYTES == DYNAMIC_SMEM_BYTES); +/// Compile-time check that every B-scale stage supports aligned 16-byte vector loads. const _: () = assert!(SFB_OFFSET.is_multiple_of(16) && SCALE_BYTES.is_multiple_of(16)); /// One lane's scale words for four pairs of B fragments. struct BScaleFragments { + /// Four pairs of packed u32 scale words, with low/high halves selecting the K64 half. words: [[u32; 2]; 4], } @@ -70,16 +90,19 @@ impl BScaleFragments { } } +/// 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, @@ -92,6 +115,7 @@ const fn output_fragment_row( 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, @@ -111,9 +135,12 @@ unsafe fn store_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))] @@ -312,11 +339,18 @@ pub mod kernels { #[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([u32; 128 * STAGES]); + 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| { @@ -324,6 +358,7 @@ mod tests { })) } + /// 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(); @@ -354,6 +389,7 @@ mod tests { } } + /// 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(); @@ -365,6 +401,7 @@ mod tests { 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 { diff --git a/cuda-oxide/cuteir/examples/blockscale_gemm_ctm/src/main.rs b/cuda-oxide/cuteir/examples/blockscale_gemm_ctm/src/main.rs index 370b75f981..7c5c700aa2 100644 --- a/cuda-oxide/cuteir/examples/blockscale_gemm_ctm/src/main.rs +++ b/cuda-oxide/cuteir/examples/blockscale_gemm_ctm/src/main.rs @@ -2,6 +2,8 @@ * 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; @@ -17,6 +19,7 @@ use kernel::{DYNAMIC_SMEM_BYTES, K, K_STAGES, KB, M, N, THREADS, TOTAL_TILES, ke /// 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 { @@ -43,6 +46,7 @@ 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 } @@ -53,10 +57,12 @@ 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) } @@ -95,6 +101,7 @@ fn pack_canonical_scales(dense: &[u8], rows: usize) -> Vec { .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); @@ -104,6 +111,7 @@ fn percentile(sorted: &[f32], numerator: usize, denominator: usize) -> f32 { 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; diff --git a/cuda-oxide/cuteir/examples/blockscale_gemm_ctm/src/physical.rs b/cuda-oxide/cuteir/examples/blockscale_gemm_ctm/src/physical.rs index 19709b4ce7..bddc6f1c1c 100644 --- a/cuda-oxide/cuteir/examples/blockscale_gemm_ctm/src/physical.rs +++ b/cuda-oxide/cuteir/examples/blockscale_gemm_ctm/src/physical.rs @@ -2,13 +2,16 @@ // 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 @@ -47,10 +50,14 @@ pub const fn output_row(warp: usize, lane: usize, band: usize, slot: usize) -> u #[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)) @@ -96,6 +103,7 @@ mod tests { } } + /// 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]; @@ -110,6 +118,7 @@ mod tests { 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 { @@ -126,6 +135,7 @@ mod tests { 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]; diff --git a/cuda-oxide/cuteir/examples/blockscale_gemm_cute/src/main.rs b/cuda-oxide/cuteir/examples/blockscale_gemm_cute/src/main.rs index 4a2167c7b0..341a562481 100644 --- a/cuda-oxide/cuteir/examples/blockscale_gemm_cute/src/main.rs +++ b/cuda-oxide/cuteir/examples/blockscale_gemm_cute/src/main.rs @@ -71,47 +71,80 @@ use cute_rs::{ 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 = L, C<32>>, T2, C<1>>>; +/// Element-domain swizzle for the FP16 carriers consumed by matrix loads. type SmemP = Composed, 0, TileP>; +/// The same packed input stage expressed as 128x64 bytes for TMA. type ByteSmem = Composed, 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 = TmaLoadPipeline; +/// TMA store queue allowing one outstanding committed output group. type EpiloguePipeline = TmaStorePipeline<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); @@ -133,8 +166,10 @@ const _: () = { 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 @@ -360,6 +395,7 @@ mod kernels { /// 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 { @@ -386,6 +422,7 @@ 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 } @@ -396,10 +433,12 @@ 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) } @@ -439,6 +478,7 @@ fn pack_canonical_scales(dense: &[u8], rows: usize) -> Vec { .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); @@ -448,6 +488,7 @@ fn percentile(sorted: &[f32], numerator: usize, denominator: usize) -> f32 { 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"); @@ -633,7 +674,9 @@ fn main() { // 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 diff --git a/cuda-oxide/cuteir/examples/elementwise_cute/README.md b/cuda-oxide/cuteir/examples/elementwise_cute/README.md index 0d122c4a31..89c4d33bd7 100644 --- a/cuda-oxide/cuteir/examples/elementwise_cute/README.md +++ b/cuda-oxide/cuteir/examples/elementwise_cute/README.md @@ -79,7 +79,7 @@ eight values per tile. ## Run it -You need the normal CUDA Oxide toolchain and an NVIDIA GPU supported by the +You need the normal cuda-oxide toolchain and an NVIDIA GPU supported by the project. From the `cuda-oxide/` directory: ```bash diff --git a/cuda-oxide/cuteir/examples/elementwise_cute/src/main.rs b/cuda-oxide/cuteir/examples/elementwise_cute/src/main.rs index aad1ae75e0..a09a10dfcd 100644 --- a/cuda-oxide/cuteir/examples/elementwise_cute/src/main.rs +++ b/cuda-oxide/cuteir/examples/elementwise_cute/src/main.rs @@ -26,14 +26,19 @@ use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; /// 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(); @@ -59,6 +64,7 @@ mod kernels { } } + /// 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(); @@ -84,10 +90,12 @@ mod kernels { } } +/// 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"); diff --git a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/Cargo.lock b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/Cargo.lock index 661f5b8b98..144042863e 100644 --- a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/Cargo.lock +++ b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/Cargo.lock @@ -251,6 +251,7 @@ dependencies = [ name = "fmha_decode_sm120a_cute" version = "0.1.0" dependencies = [ + "ctm-rs", "cuda-core", "cuda-device", "cuda-host", diff --git a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/Cargo.toml b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/Cargo.toml index 277691e0ec..1879dcad0a 100644 --- a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/Cargo.toml +++ b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/Cargo.toml @@ -9,6 +9,8 @@ description = "SM120a: Dedicated split-KV FP16 attention decode with typed CuTe [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"] } diff --git a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/README.md b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/README.md index 2ae7176a2a..cf55a75b4e 100644 --- a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/README.md +++ b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/README.md @@ -4,6 +4,12 @@ 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. +The kernel uses `cute_rs` for typed layouts and copies, and `ctm_rs` for +instruction-level operations. CTM also exposes the kernel macros and canonical +storage types. The manifest retains `cuda-device` because the macro expansions +still refer to that crate internally. Dot products and split reduction retain +their scalar FP32 arithmetic. + - Batch size **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), @@ -86,3 +92,4 @@ use the same method and sampling protocol to assess a change. See - [CuTeDSL reference and setup](reference/README.md) - [Validation, benchmark commands, and detailed results](reference/measurement-notes.md) - [Latest compiler regression checks](reference/compiler-regression-results.json) · [Earlier validation](reference/validation-results.json) +- [CTM migration: correctness, instruction equivalence, and matched timings](reference/ctm-migration-results.json) diff --git a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/reference/ctm-migration-results.json b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/reference/ctm-migration-results.json new file mode 100644 index 0000000000..5b17a11e3b --- /dev/null +++ b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/reference/ctm-migration-results.json @@ -0,0 +1,283 @@ +{ + "status": "passed", + "date": "2026-10-14", + "change": "Route SM120a attention GPU primitives through CTM-RS, preserving CuTe layouts and copies.", + "base_revision": "238a2ef910fc312fa5352645234a16ef333c9930", + "target": "sm_120a", + "device_backend": "cutlass-mlir", + "nvvm_dialect_unchanged": true, + "environment": { + "gpu": "0, NVIDIA GeForce RTX 5090, GPU-04a6d226-f1d6-0fc8-a95a-3ede21981e72, 00000000:02:00.0, 615.71.09, 12.0", + "cuda_toolkit": "13.4.92", + "cutlass": "4.7.0", + "nsight_systems": "2026.3.2.476", + "compute_sanitizer": "2026.3.0.0", + "rustc": "1.101.0-nightly (6f5da406b 2026-10-10)", + "cutlass_compiler_sha256": "57df017e3585a10443c74c8b4cd99bda854242fb2f4c9534cf56d58c2c741628" + }, + "source_sha256": { + "cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/kernel.rs": "4173e3c691b7e0e9412ef434a2d2601bd76d1d076630b292ac18404b4b4ad868", + "cuda-oxide/cuteir/cute-rs/src/mma.rs": "a180d77d76b2598f4f784300c9f776aa55c67db06e8ca64b337d168b44ef2674", + "cuda-oxide/cuteir/ctm-rs/src/lib.rs": "c48611a7c50b36b6f2b10f07d6d01ae69951692ec35d3dd8da1f71f37e3b41e9", + "cuda-oxide/cuteir/ctm-rs/src/memory.rs": "3ede95ab877c0d99d38ce509d32d31dc9690a8203907c5a96ebb240ebb340b00", + "cuda-oxide/cuteir/ctm-rs/src/async_copy.rs": "f7599dd7e5105c0ab9f4f6c4fb9d6a95c668287b15f9729a7b9f2e8b375b06de", + "cuda-oxide/cuteir/ctm-rs/src/float.rs": "1f4ec19f491d69d86e6dc24b226524ab3553a4d2577673f79f9ad93921c469a1", + "cuda-oxide/cuteir/ctm-rs/src/convert.rs": "69cb75675119cea2acc5e277eff55dab792012e9ada03b36720c1d57d1192e30", + "cuda-oxide/cuteir/ctm-rs/src/warp.rs": "40e066cba7cc253f065f5857c3de6b74009ed04439660f13e0cc6dcfb1ece3d1", + "cuda-oxide/cuteir/ctm-rs/src/mma.rs": "b37b6fbdc37a43dddcfd3a1180370f88c7e6820d8f4776d34d06460a27b0f95a" + }, + "build_command": "CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir cargo oxide build fmha_decode_sm120a_cute --arch sm_120a", + "build_working_directory": "cuda-oxide", + "cubins": { + "archived": { + "record": "compiler-regression-results.json", + "sha256": "e64a7a8ebf459d0d91640541660e8c0870cf83c64a6318e6f22ca417f7df0d43" + }, + "before_sha256": "19faa24069d5cbd4474c925afb8c8daea03218b82fea16aaf35931840b08deff", + "after_sha256": "2c0a33147ca25ca9ed456abfb461685f0dbd74e2c1085992032707147825e7d0" + }, + "codegen": { + "all_kernel_instruction_bytes_equal_to_before_and_archive": true, + "all_kernel_resources_equal": true, + "comparison": "Complete kernel instruction encodings, including padding, checked against actual ELF code sections; symbol hashes and non-code metadata may differ.", + "kernels": [ + { + "symbol": "fmha_decode_merge_TID_5e269d3319f1e427553827620d3f54d1", + "instructions": 424, + "instruction_sha256": "3eb994c9d5888372f0cbe06a9401e43498c23e1ee41fa62cd1983564e341f7a4", + "resources": "REG:57 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:936 TEXTURE:0 SURFACE:0 SAMPLER:0" + }, + { + "symbol": "fmha_decode_partial_TID_c77f7dcbd686a8c169da40275274c297", + "instructions": 584, + "instruction_sha256": "a7f974504561047f89ca7be003d9bb340f1df131915adc25bfc243c252ee3c03", + "resources": "REG:40 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:968 TEXTURE:0 SURFACE:0 SAMPLER:0" + }, + { + "symbol": "fmha_decode_partial_TID_d56bd9bee4dd45adb75f06f78b307df2", + "instructions": 616, + "instruction_sha256": "bac389ccc0ce875475d3aa3f5fb77cb77640abfbdcccbbee6542bfd9627e4c7c", + "resources": "REG:40 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:968 TEXTURE:0 SURFACE:0 SAMPLER:0" + }, + { + "symbol": "fmha_decode_partial_TID_315f39bba02f87e27b79ac1dff6dfe68", + "instructions": 664, + "instruction_sha256": "e50dd73002a9bd42c37e48bdf92a733773a2668790ab5160dc16bcbd846c2fb8", + "resources": "REG:40 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:968 TEXTURE:0 SURFACE:0 SAMPLER:0" + }, + { + "symbol": "fmha_decode_merge_TID_41adcf84b12795e46849177416d989a2", + "instructions": 272, + "instruction_sha256": "bd9c688db922209d7ce5b9a7836296d3f5bb2900a38fec0cc101e76944a3445e", + "resources": "REG:76 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:936 TEXTURE:0 SURFACE:0 SAMPLER:0" + }, + { + "symbol": "fmha_decode_merge_TID_fc33f53cb6773dead155a9ffbf71a0c4", + "instructions": 64, + "instruction_sha256": "4dbae85d1a32a5fd460ee3e3c6620c9db400a1ebb2af0fb7597d6be592f598b7", + "resources": "REG:18 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:936 TEXTURE:0 SURFACE:0 SAMPLER:0" + }, + { + "symbol": "fmha_decode_merge_TID_5518c6526b4f33c460240fafbe312180", + "instructions": 48, + "instruction_sha256": "e3629f757a9adfe12f504bddc381cbb796802bc2840150c84d68301d770beb37", + "resources": "REG:16 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:936 TEXTURE:0 SURFACE:0 SAMPLER:0" + }, + { + "symbol": "fmha_decode_merge_TID_c3ff65afab464634a027b55f4b1b0b20", + "instructions": 304, + "instruction_sha256": "b339fd105a10af9c1ace89df8a3d15e5250530ec53e7035f42e6f6caa54734e5", + "resources": "REG:48 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:936 TEXTURE:0 SURFACE:0 SAMPLER:0" + }, + { + "symbol": "fmha_decode_partial_TID_c265f8725d5316f759702aa036961cc4", + "instructions": 608, + "instruction_sha256": "a328b731b83029b3cf8abe31dad5a523ef5c524ce528b6dbb7f62c2c89adf4cc", + "resources": "REG:40 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:968 TEXTURE:0 SURFACE:0 SAMPLER:0" + }, + { + "symbol": "fmha_decode_partial_TID_ec17509d1a88420e81cbd4fcf5f3c76e", + "instructions": 568, + "instruction_sha256": "c3ee6cb0028986d0150681098705649bc58a46ad2a73737d8ccb1a3398971bff", + "resources": "REG:40 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:968 TEXTURE:0 SURFACE:0 SAMPLER:0" + } + ] + }, + "correctness": { + "oracle": "Every output checked against independent CPU FP64 stable softmax; poisoned output buffers and repeat-launch checks.", + "suites": [ + { + "suite": "decode-validation", + "case_count": 52, + "all_outputs_bitwise_identical": true + } + ], + "decode_runtime_static_cross_checks": 6, + "cutedsl_benchmark_input_hash_and_output_parity": "passed in every benchmark case" + }, + "sanitizers": { + "runs": 8, + "all_passed": true, + "tools": [ + "memcheck", + "racecheck", + "synccheck", + "initcheck (global and shared)" + ], + "cases": [ + "runtime: Sk=257,H=3,input-scale=16", + "static: Sk=65,H=3,input-scale=16" + ] + }, + "timing": { + "protocol": "Nsight Systems, 100 warmups, 31 samples per implementation per round, 3 alternating rounds; median of round medians", + "shape": { + "batch": 1, + "seq_q": 1, + "seq_k": 1024, + "heads": 8, + "head_dim": 64 + }, + "cases": [ + { + "case": "decode-runtime-noncausal", + "before_report": "/tmp/ctm-sm120a-validation/before/benchmark-decode-runtime-noncausal/decode-8rcg5a4e/report.json", + "after_report": "/tmp/ctm-sm120a-validation/after/benchmark-decode-runtime-noncausal/decode-kpgam06t/report.json", + "round_medians_us": { + "active": { + "before_rust": [ + 2.944, + 2.944, + 2.944 + ], + "after_rust": [ + 2.944, + 2.944, + 2.944 + ], + "before_reference": [ + 3.072, + 3.072, + 3.072 + ], + "after_reference": [ + 3.072, + 3.103, + 3.104 + ] + }, + "wall": { + "before_rust": [ + 3.232, + 3.2, + 3.2 + ], + "after_rust": [ + 3.232, + 3.2, + 3.264 + ], + "before_reference": [ + 3.744, + 3.616, + 3.648 + ], + "after_reference": [ + 3.552, + 3.936, + 3.584 + ] + } + }, + "rust_before_active_us": 2.944, + "rust_after_active_us": 2.944, + "rust_active_delta_percent": 0.0, + "reference_before_active_us": 3.072, + "reference_after_active_us": 3.103, + "rust_before_wall_us": 3.2, + "rust_after_wall_us": 3.232, + "rust_wall_delta_percent": 1.0000000000000009, + "reference_before_wall_us": 3.648, + "reference_after_wall_us": 3.584 + }, + { + "case": "decode-static-noncausal", + "before_report": "/tmp/ctm-sm120a-validation/before/benchmark-decode-static-noncausal/decode-vq3y72n4/report.json", + "after_report": "/tmp/ctm-sm120a-validation/after/benchmark-decode-static-noncausal/decode-ot6rsa_n/report.json", + "round_medians_us": { + "active": { + "before_rust": [ + 2.432, + 2.432, + 2.432 + ], + "after_rust": [ + 2.432, + 2.432, + 2.432 + ], + "before_reference": [ + 2.208, + 2.208, + 2.208 + ], + "after_reference": [ + 2.208, + 2.24, + 2.24 + ] + }, + "wall": { + "before_rust": [ + 2.752, + 2.752, + 2.816 + ], + "after_rust": [ + 2.72, + 2.72, + 2.784 + ], + "before_reference": [ + 2.912, + 3.072, + 3.2 + ], + "after_reference": [ + 2.816, + 3.008, + 2.848 + ] + } + }, + "rust_before_active_us": 2.432, + "rust_after_active_us": 2.432, + "rust_active_delta_percent": 0.0, + "reference_before_active_us": 2.208, + "reference_after_active_us": 2.24, + "rust_before_wall_us": 2.752, + "rust_after_wall_us": 2.72, + "rust_wall_delta_percent": -1.1627906976743985, + "reference_before_wall_us": 3.072, + "reference_after_wall_us": 2.848 + } + ], + "interpretation": "Kernel-active before/after differences are observed run variation, with identical generated code. Decode active sum and sequence span are separate metrics.", + "archived_comparison": "Archived timings used physical GPU 1; this matched before/after pair used physical GPU 0. Compare the paired measurements to assess the migration, not absolute timing across boards." + }, + "host_checks": { + "ctm_rs_tests_including_doctests": 20, + "cute_rs_tests_including_doctests": 63, + "example_tests": 19, + "python_tests_run": 28, + "python_result": "OK", + "clippy_warnings_denied": "passed", + "rustdoc_warnings_denied": "passed", + "primitives_api_compile": "passed; nine exact instruction identities, selectors, operands and result adaptation checked" + }, + "raw_evidence": "/tmp/ctm-sm120a-validation", + "instruction_comparison": "/tmp/ctm-sm120a-validation/after/fmha_decode_sm120a_cute.comparison.json", + "numerical_output_hashes": "/tmp/ctm-sm120a-validation/output-correctness-comparison.json", + "sanitizer_commands_logs": "/tmp/ctm-sm120a-validation/after/correctness-sanitizers-summary.json", + "primitive_compile_manifest": "/tmp/ctm-sm120a-validation/primitives/manifest.json" +} diff --git a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/config.rs b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/config.rs index 02babacb19..c3ad013501 100644 --- a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/config.rs +++ b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/config.rs @@ -3,23 +3,35 @@ * 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", @@ -29,30 +41,48 @@ impl Specialization { } } +/// 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()); @@ -92,6 +122,7 @@ impl LaunchPlan { } 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)) { @@ -104,6 +135,7 @@ impl Args { } } + /// 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)), @@ -112,6 +144,7 @@ impl Args { } } + /// 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, @@ -214,6 +247,7 @@ impl Args { } } +/// Reject a static kernel whose dimensions differ from the workload. pub fn validate_kernel_shape( shape: Shape, ) -> Result<(), String> { @@ -229,10 +263,12 @@ pub fn validate_kernel_shape( 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, @@ -242,10 +278,12 @@ mod tests { } } + /// 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()); @@ -254,6 +292,7 @@ mod tests { 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)] { @@ -275,6 +314,7 @@ mod tests { } } + /// 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)] { @@ -294,6 +334,7 @@ mod tests { } } + /// Reject options that cannot participate in the decode comparison matrix. #[test] fn all_cases_rejects_standalone_execution_output_and_causal_attention() { for cli in [ @@ -306,6 +347,7 @@ mod tests { } } + /// 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 [ @@ -328,6 +370,7 @@ mod tests { } } + /// Reject empty dimensions, multiple queries, and causal decode. #[test] fn launch_plan_rejects_invalid_decode_workloads() { for invalid in [ @@ -350,6 +393,7 @@ mod tests { } } + /// Check both tensor and scratch indexing limits before allocation. #[test] fn launch_plan_rejects_input_and_scratch_index_overflow() { for invalid in [ @@ -365,6 +409,7 @@ mod tests { 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(); @@ -414,6 +459,7 @@ mod tests { 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 [ @@ -435,6 +481,7 @@ mod tests { } } + /// Check that every supported shape selects the requested static or runtime kernel. #[test] fn specialization_selects_only_exact_aot_shapes() { assert_eq!( @@ -463,6 +510,7 @@ mod tests { } } + /// 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)] { @@ -496,6 +544,7 @@ mod tests { } } + /// 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()); diff --git a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/kernel.rs b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/kernel.rs index 9e28b17383..47434e52da 100644 --- a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/kernel.rs +++ b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/kernel.rs @@ -9,20 +9,28 @@ //! 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 cuda_device::{ - DisjointSlice, DynamicSharedArray, async_copy, barrier, cuda_module, float, kernel, - launch_bounds, launch_contract, thread, warp, +use ctm_rs::{ + async_copy, cuda_module, float, kernel, launch_bounds, launch_contract, mbarrier, + memory::{DisjointSlice, DynamicSharedArray}, + thread, warp, }; use cute_rs::markers::{C, Composed, CpAsync, L, LeadingDim, RowMajor, Swizzle, T2}; use cute_rs::{GmemMatrix, SmemTile, copy_g2s}; +/// Swizzled 64-by-64 FP16 layout used for both K and V in shared memory. type KvLayout = Composed, 0, RowMajor<64, 64>>; +/// Swizzled 16-by-64 region covered by one cooperative copy. type CopyLayout = Composed, 0, RowMajor<16, 64>>; +/// Map 128 threads to 16 rows and eight 16-byte copy positions. type CopyThreads = L, C<8>>, T2, C<1>>>; +/// Eight adjacent FP16 values copied by each thread. type CopyValues = L, C<8>>, T2, 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 @@ -33,17 +41,20 @@ const fn assert_shape() { 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) => { - $({ const $index: usize = $value; $body })* + $({ #[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] { @@ -52,6 +63,7 @@ fn warp_max(mut value: f32) -> f32 { 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] { @@ -60,12 +72,17 @@ fn warp_sum(mut value: f32) -> f32 { 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 { @@ -75,6 +92,10 @@ unsafe fn tile(base: *mut f16) -> KvTile { } } +/// 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] { @@ -90,10 +111,12 @@ unsafe fn stage(src: &GmemMatrix>, split: usize, dst: &KvTile }); } +/// 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))] @@ -142,8 +165,8 @@ pub mod kernels { 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); + async_copy::cp_async_wait_group::<0>(); + mbarrier::barrier_cta_sync(1, 128); // 32 independent keys at a time, with four lanes per 64-D dot. static_for!(PASS in [0, 1] { @@ -162,7 +185,7 @@ pub mod kernels { } else { f32::NEG_INFINITY }; } }); - barrier::barrier_cta_sync(1, 128); + mbarrier::barrier_cta_sync(1, 128); let score = if tid < 64 { *scores.add(tid) } else { @@ -172,7 +195,7 @@ pub mod kernels { if lane == 0 { *maxima.add(warp_id) = local_maximum; } - barrier::barrier_cta_sync(1, 128); + mbarrier::barrier_cta_sync(1, 128); let maximum = max( max(*maxima, *maxima.add(1)), max(*maxima.add(2), *maxima.add(3)), @@ -186,7 +209,7 @@ pub mod kernels { if lane == 0 { *sums.add(warp_id) = local_sum; } - barrier::barrier_cta_sync(1, 128); + mbarrier::barrier_cta_sync(1, 128); // Each warp integrates 16 keys into two output dimensions per lane. let mut numerator = [0.0f32; 2]; @@ -200,7 +223,7 @@ pub mod kernels { static_for!(D in [0, 1] { *warp_outputs.add(warp_id * 64 + lane * 2 + D) = numerator[D]; }); - barrier::barrier_cta_sync(1, 128); + mbarrier::barrier_cta_sync(1, 128); if warp_id == 0 { let dst = partials .as_mut_ptr() @@ -221,6 +244,7 @@ pub mod kernels { } } + /// 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] diff --git a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/main.rs b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/main.rs index d9dd45222b..021a259b7a 100644 --- a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/main.rs +++ b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/main.rs @@ -3,16 +3,21 @@ * 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/cuteir/examples/fmha_decode_sm120a_cute/src/runner.rs b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/runner.rs index 105a412161..7254d32b2b 100644 --- a/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/runner.rs +++ b/cuda-oxide/cuteir/examples/fmha_decode_sm120a_cute/src/runner.rs @@ -17,8 +17,10 @@ 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!( @@ -66,6 +68,7 @@ pub fn run() -> Result<()> { } } +/// 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)?; diff --git a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/Cargo.lock b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/Cargo.lock index f76b544bbb..61c247deb4 100644 --- a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/Cargo.lock +++ b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/Cargo.lock @@ -251,6 +251,7 @@ dependencies = [ name = "fmha_prefill_sm120a_cute" version = "0.1.0" dependencies = [ + "ctm-rs", "cuda-core", "cuda-device", "cuda-host", diff --git a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/Cargo.toml b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/Cargo.toml index 2feef6f2a5..52dfb10e64 100644 --- a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/Cargo.toml +++ b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/Cargo.toml @@ -9,6 +9,8 @@ description = "SM120a: FP16 prefill and FMHA with warp MMA and online softmax in [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"] } diff --git a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/README.md b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/README.md index ccb8931922..9067822b5c 100644 --- a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/README.md +++ b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/README.md @@ -3,6 +3,12 @@ This cuda-oxide example computes `softmax(Q @ K.T / sqrt(64)) @ V` using FP16 inputs, FP32 accumulation, and 64-by-64 tiles. +The kernel uses `cute_rs` for typed layouts, copies, and matrix fragments, +and `ctm_rs` for instruction-level operations, including the FP16 MMA behind +CuTe's accumulator wrapper. CTM also exposes the kernel macros and canonical +storage types. The manifest retains `cuda-device` because the macro expansions +still refer to that crate internally. + - Batch size **1**, head dimension **64**, contiguous BSHD tensors. - Sequence lengths are positive multiples of 64. **Causal** (default) requires equal query/KV lengths; `--noncausal` permits unequal lengths and also `Sq=1`. @@ -85,4 +91,5 @@ work. This is an equal-weight summary; full parity is not yet established. - [CuTeDSL reference and setup](reference/README.md) - [Validation, benchmark commands, and detailed results](reference/measurement-notes.md) - [Profiling findings and latest validation](reference/profiling-notes.md) +- [CTM migration: correctness, instruction equivalence, and matched timings](reference/ctm-migration-results.json) - [Dedicated decode example](../fmha_decode_sm120a_cute/README.md) diff --git a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/reference/ctm-migration-results.json b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/reference/ctm-migration-results.json new file mode 100644 index 0000000000..db594fac6a --- /dev/null +++ b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/reference/ctm-migration-results.json @@ -0,0 +1,258 @@ +{ + "status": "passed", + "date": "2026-10-14", + "change": "Route SM120a attention GPU primitives through CTM-RS, preserving CuTe layouts and copies.", + "base_revision": "238a2ef910fc312fa5352645234a16ef333c9930", + "target": "sm_120a", + "device_backend": "cutlass-mlir", + "nvvm_dialect_unchanged": true, + "environment": { + "gpu": "0, NVIDIA GeForce RTX 5090, GPU-04a6d226-f1d6-0fc8-a95a-3ede21981e72, 00000000:02:00.0, 615.71.09, 12.0", + "cuda_toolkit": "13.4.92", + "cutlass": "4.7.0", + "nsight_systems": "2026.3.2.476", + "compute_sanitizer": "2026.3.0.0", + "rustc": "1.101.0-nightly (6f5da406b 2026-10-10)", + "cutlass_compiler_sha256": "57df017e3585a10443c74c8b4cd99bda854242fb2f4c9534cf56d58c2c741628" + }, + "source_sha256": { + "cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/kernel.rs": "f0d56299428a846d778d9057eba11acb1f8aed64e242462dfe67d1143dba1db6", + "cuda-oxide/cuteir/cute-rs/src/mma.rs": "a180d77d76b2598f4f784300c9f776aa55c67db06e8ca64b337d168b44ef2674", + "cuda-oxide/cuteir/ctm-rs/src/lib.rs": "c48611a7c50b36b6f2b10f07d6d01ae69951692ec35d3dd8da1f71f37e3b41e9", + "cuda-oxide/cuteir/ctm-rs/src/memory.rs": "3ede95ab877c0d99d38ce509d32d31dc9690a8203907c5a96ebb240ebb340b00", + "cuda-oxide/cuteir/ctm-rs/src/async_copy.rs": "f7599dd7e5105c0ab9f4f6c4fb9d6a95c668287b15f9729a7b9f2e8b375b06de", + "cuda-oxide/cuteir/ctm-rs/src/float.rs": "1f4ec19f491d69d86e6dc24b226524ab3553a4d2577673f79f9ad93921c469a1", + "cuda-oxide/cuteir/ctm-rs/src/convert.rs": "69cb75675119cea2acc5e277eff55dab792012e9ada03b36720c1d57d1192e30", + "cuda-oxide/cuteir/ctm-rs/src/warp.rs": "40e066cba7cc253f065f5857c3de6b74009ed04439660f13e0cc6dcfb1ece3d1", + "cuda-oxide/cuteir/ctm-rs/src/mma.rs": "b37b6fbdc37a43dddcfd3a1180370f88c7e6820d8f4776d34d06460a27b0f95a" + }, + "build_command": "CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir cargo oxide build fmha_prefill_sm120a_cute --arch sm_120a", + "build_working_directory": "cuda-oxide", + "cubins": { + "archived": { + "record": "epilogue-results.json", + "sha256": "3ad0811fc76a3671a002ad351d6b2ec484ad0429dc48b79007ed9141e5113f81" + }, + "before_sha256": "4f7c2f591edb7216d451b7775ba5fb37f03373331cb9656d5b1532e5db08d20c", + "after_sha256": "3dc4ba4a1b4ee5d235ca99c8efcbad2e4c01dcceec6358d17268d1cf938197ac" + }, + "codegen": { + "all_kernel_instruction_bytes_equal_to_before_and_archive": true, + "all_kernel_resources_equal": true, + "comparison": "Complete kernel instruction encodings, including padding, checked against actual ELF code sections; symbol hashes and non-code metadata may differ.", + "kernels": [ + { + "symbol": "fmha_prefill_TID_33beb4c1a569d292f6a1d474c4c30ff6", + "instructions": 888, + "instruction_sha256": "6629648ac76f9875866f9d3961b366263802c76380b1b75f274582c01227723e", + "resources": "REG:120 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:972 TEXTURE:0 SURFACE:0 SAMPLER:0" + }, + { + "symbol": "fmha_prefill_TID_3982e4cfaa04d93295cffa97b4941569", + "instructions": 1136, + "instruction_sha256": "5ae1603c6e09b55b29973d6becbc78593266f7ce9ec4a2028c0d197f80d91c2b", + "resources": "REG:128 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:972 TEXTURE:0 SURFACE:0 SAMPLER:0" + }, + { + "symbol": "fmha_prefill_TID_ab065f1ad81b2a331a0ce50d6ab568f2", + "instructions": 1056, + "instruction_sha256": "8e8cb15828544e9b735c283c02bb9455a0c105344d5c173694e9a2c4da9f0872", + "resources": "REG:136 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:972 TEXTURE:0 SURFACE:0 SAMPLER:0" + }, + { + "symbol": "fmha_prefill_TID_6c09702c4d47a1647a97f344eaaf5855", + "instructions": 984, + "instruction_sha256": "5c8bf499ce0e19c73c9d69b41c827407b340b82a0af921ec32249e0cf15760a3", + "resources": "REG:125 STACK:0 SHARED:1024 LOCAL:0 CONSTANT[0]:972 TEXTURE:0 SURFACE:0 SAMPLER:0" + } + ] + }, + "correctness": { + "oracle": "Every output checked against independent CPU FP64 stable softmax; poisoned output buffers and repeat-launch checks.", + "suites": [ + { + "suite": "prefill-validation-runtime", + "case_count": 15, + "all_outputs_bitwise_identical": true + }, + { + "suite": "prefill-validation-static", + "case_count": 4, + "all_outputs_bitwise_identical": true + } + ], + "decode_runtime_static_cross_checks": 0, + "cutedsl_benchmark_input_hash_and_output_parity": "passed in every benchmark case" + }, + "sanitizers": { + "runs": 16, + "all_passed": true, + "tools": [ + "memcheck", + "racecheck", + "synccheck", + "initcheck (global and shared)" + ], + "cases": [ + "runtime causal/noncausal: Sq=Sk=128,H=3", + "static causal/noncausal: Sq=Sk=1024,H=8" + ] + }, + "timing": { + "protocol": "Nsight Systems, 100 warmups, 31 samples per implementation per round, 3 alternating rounds; median of round medians", + "shape": { + "batch": 1, + "seq_q": 1024, + "seq_k": 1024, + "heads": 8, + "head_dim": 64 + }, + "cases": [ + { + "case": "prefill-runtime-causal", + "before_report": "/tmp/ctm-sm120a-validation/before/benchmark-prefill-runtime-causal/run-iih15ir_/report.json", + "after_report": "/tmp/ctm-sm120a-validation/after/benchmark-prefill-runtime-causal/run-dxva2mq0/report.json", + "round_medians_us": { + "before_rust": [ + 16.064, + 16.064, + 16.063 + ], + "after_rust": [ + 16.064, + 16.063, + 16.064 + ], + "before_reference": [ + 15.871, + 15.84, + 15.84 + ], + "after_reference": [ + 15.936, + 15.935, + 15.904 + ] + }, + "rust_before_active_us": 16.064, + "rust_after_active_us": 16.064, + "rust_active_delta_percent": 0.0, + "reference_before_active_us": 15.84, + "reference_after_active_us": 15.935 + }, + { + "case": "prefill-runtime-noncausal", + "before_report": "/tmp/ctm-sm120a-validation/before/benchmark-prefill-runtime-noncausal/run-vaw1acvr/report.json", + "after_report": "/tmp/ctm-sm120a-validation/after/benchmark-prefill-runtime-noncausal/run-bc_fxe_b/report.json", + "round_medians_us": { + "before_rust": [ + 16.32, + 16.32, + 16.32 + ], + "after_rust": [ + 16.32, + 16.32, + 16.32 + ], + "before_reference": [ + 15.872, + 15.872, + 15.872 + ], + "after_reference": [ + 15.904, + 15.936, + 15.936 + ] + }, + "rust_before_active_us": 16.32, + "rust_after_active_us": 16.32, + "rust_active_delta_percent": 0.0, + "reference_before_active_us": 15.872, + "reference_after_active_us": 15.936 + }, + { + "case": "prefill-static-causal", + "before_report": "/tmp/ctm-sm120a-validation/before/benchmark-prefill-static-causal/run-ussrwr19/report.json", + "after_report": "/tmp/ctm-sm120a-validation/after/benchmark-prefill-static-causal/run-aib0670m/report.json", + "round_medians_us": { + "before_rust": [ + 15.68, + 15.68, + 15.68 + ], + "after_rust": [ + 15.712, + 15.744, + 15.744 + ], + "before_reference": [ + 15.616, + 15.648, + 15.616 + ], + "after_reference": [ + 15.711, + 15.712, + 15.712 + ] + }, + "rust_before_active_us": 15.68, + "rust_after_active_us": 15.744, + "rust_active_delta_percent": 0.40816326530612734, + "reference_before_active_us": 15.616, + "reference_after_active_us": 15.712 + }, + { + "case": "prefill-static-noncausal", + "before_report": "/tmp/ctm-sm120a-validation/before/benchmark-prefill-static-noncausal/run-89klqn7u/report.json", + "after_report": "/tmp/ctm-sm120a-validation/after/benchmark-prefill-static-noncausal/run-xppeul6e/report.json", + "round_medians_us": { + "before_rust": [ + 15.776, + 15.776, + 15.807 + ], + "after_rust": [ + 15.84, + 15.84, + 15.84 + ], + "before_reference": [ + 15.648, + 15.648, + 15.648 + ], + "after_reference": [ + 15.712, + 15.712, + 15.712 + ] + }, + "rust_before_active_us": 15.776, + "rust_after_active_us": 15.84, + "rust_active_delta_percent": 0.40567951318457585, + "reference_before_active_us": 15.648, + "reference_after_active_us": 15.712 + } + ], + "interpretation": "Kernel-active before/after differences are observed run variation, with identical generated code. Decode active sum and sequence span are separate metrics.", + "archived_comparison": "Archived timings used physical GPU 1; this matched before/after pair used physical GPU 0. Compare the paired measurements to assess the migration, not absolute timing across boards." + }, + "host_checks": { + "ctm_rs_tests_including_doctests": 20, + "cute_rs_tests_including_doctests": 63, + "example_tests": 14, + "python_tests_run": 29, + "python_result": "OK (skipped=1)", + "clippy_warnings_denied": "passed", + "rustdoc_warnings_denied": "passed", + "primitives_api_compile": "passed; nine exact instruction identities, selectors, operands and result adaptation checked" + }, + "raw_evidence": "/tmp/ctm-sm120a-validation", + "instruction_comparison": "/tmp/ctm-sm120a-validation/after/fmha_prefill_sm120a_cute.comparison.json", + "numerical_output_hashes": "/tmp/ctm-sm120a-validation/output-correctness-comparison.json", + "sanitizer_commands_logs": "/tmp/ctm-sm120a-validation/after/correctness-sanitizers-summary.json", + "primitive_compile_manifest": "/tmp/ctm-sm120a-validation/primitives/manifest.json" +} diff --git a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/comparison.rs b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/comparison.rs index 38b92afb55..47fe60a373 100644 --- a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/comparison.rs +++ b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/comparison.rs @@ -3,6 +3,8 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Shared process launcher for the SM120a prefill and decode comparisons. + use crate::fixture::Shape; use std::{ error::Error, @@ -10,8 +12,10 @@ use std::{ 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>, @@ -32,6 +36,7 @@ fn resolve_python( }) } +/// Reject Python selection in standalone mode and raw output in comparison mode. pub fn validate_mode( rust_only: bool, python: Option<&PathBuf>, @@ -46,6 +51,7 @@ pub fn validate_mode( Ok(()) } +/// Run the Python comparison with the same shape, inputs, and timing options. #[allow(clippy::too_many_arguments)] pub fn run( workload: &str, @@ -136,15 +142,19 @@ pub fn run( Ok(()) } +/// Tests for interpreter overrides and virtualenv fallback order. #[cfg(test)] mod tests { use super::*; use std::sync::atomic::{AtomicUsize, Ordering}; - struct TempDirectory(PathBuf); + /// 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-{}-{}", @@ -155,6 +165,7 @@ mod tests { 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(); @@ -169,6 +180,7 @@ mod tests { } } + /// Check that explicit and environment Python choices override discovered virtualenvs. #[test] fn python_overrides_take_precedence_over_virtualenvs() { let temp = TempDirectory::new(); @@ -187,6 +199,7 @@ mod tests { ); } + /// 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(); diff --git a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/fixture.rs b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/fixture.rs index 46cc5bc993..dc327a3c5c 100644 --- a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/fixture.rs +++ b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/fixture.rs @@ -5,23 +5,36 @@ //! 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 { + /// 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()); @@ -46,15 +59,18 @@ impl Shape { 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); @@ -63,6 +79,7 @@ fn mix32(mut value: u32) -> u32 { 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| { @@ -150,9 +167,13 @@ pub fn reference(shape: Shape, q: &[f16], k: &[f16], v: &[f16]) -> Vec { 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, } @@ -190,10 +211,12 @@ pub fn verify(got: &[f16], expected: &[f64]) -> Result { 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 { @@ -231,6 +254,7 @@ mod tests { ); } + /// Check causal masking and independent head indexing using uniform logits. #[test] fn oracle_masks_future_keys_and_keeps_heads_separate() { let mut shape = Shape { @@ -257,6 +281,7 @@ mod tests { 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 { @@ -277,6 +302,7 @@ mod tests { 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]; @@ -289,6 +315,7 @@ mod tests { } } + /// Check positive dimensions, tile boundaries, and equal-length causal sequences. #[test] fn shape_rejects_zero_and_asymmetric_causal_sequences() { let good = Shape { diff --git a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/kernel.rs b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/kernel.rs index 6f0d1be789..b77c6d2b47 100644 --- a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/kernel.rs +++ b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/kernel.rs @@ -6,36 +6,48 @@ //! 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 cuda_device::{ - DisjointSlice, DynamicSharedArray, async_copy, cuda_module, float, kernel, launch_bounds, - launch_contract, thread, warp, +use ctm_rs::{ + async_copy, convert, cuda_module, float, kernel, launch_bounds, launch_contract, mbarrier, + memory::{DisjointSlice, DynamicSharedArray}, + thread, warp, }; use cute_rs::markers::{C, Composed, CpAsync, L, LeadingDim, RowMajor, Swizzle, T2}; use cute_rs::{AccC, FragA, FragBPair, GmemMatrix, SmemTile, copy_g2s}; +/// Swizzled 64-by-64 FP16 layout shared by Q, K, V, and the output staging tile. type SharedLayout = Composed, 0, RowMajor<64, 64>>; +/// Swizzled 16-by-64 subview used by one cooperative copy. type CopyLayout = Composed, 0, RowMajor<16, 64>>; +/// Map 128 threads to 16 rows and eight 16-byte copy positions. type CopyThreads = L, C<8>>, T2, C<1>>>; +/// Eight adjacent FP16 values copied by each thread. type CopyValues = L, C<8>>, T2, 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) => { - $({ const $index: usize = $value; $body })* + $({ #[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([AccC; 8]); +struct WarpTile(#[doc = "Eight adjacent 16-by-8 FP32 accumulator fragments."] [AccC; 8]); impl WarpTile { + /// Create an all-zero warp accumulator tile. #[inline(always)] fn zero() -> Self { Self([AccC::ZERO; 8]) } + /// Accumulate one K16 step across all eight column fragments. + /// + /// # Safety + /// All warp lanes must participate with matching, valid operand fragments. #[inline(always)] unsafe fn mma(&mut self, a: FragA, b: [FragBPair; 4]) { static_for!(PAIR in [0, 1, 2, 3] { @@ -68,6 +80,7 @@ impl WarpTile { 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] { @@ -92,6 +105,7 @@ impl WarpTile { pairs[0].max(pairs[1]) } + /// Mask and exponentiate scores, updating the running normalization and output scale. #[inline(always)] fn softmax( &mut self, @@ -116,6 +130,7 @@ impl WarpTile { 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] { @@ -140,6 +155,10 @@ impl WarpTile { }); } + /// 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] { @@ -161,19 +180,24 @@ impl WarpTile { } } +/// 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 { - cuda_device::convert::cvt_f16x2_f32(a, b) + convert::cvt_f16x2_f32(a, b) } /// Physical coordinate of the typed `Swizzle<3,3,3>` shared layout. @@ -184,6 +208,10 @@ fn shared_offset(row: u32, column: u32) -> usize { (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 { @@ -193,6 +221,10 @@ unsafe fn tile(base: *mut f16) -> SharedTile { } } +/// 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] { @@ -227,6 +259,10 @@ unsafe fn stage_full( 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) -> [FragBPair; 4] { let mut fragments = [FragBPair([0; 4]); 4]; @@ -236,6 +272,10 @@ unsafe fn load_k(src: &SharedTile, lane: u32) -> [FragBPair; 4] 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) -> [FragBPair; 4] { let mut fragments = [FragBPair([0; 4]); 4]; @@ -248,15 +288,25 @@ unsafe fn load_v(src: &SharedTile, lane: u32) -> [FragBPair; 4] /// 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, @@ -270,8 +320,8 @@ impl TilePipeline { // 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); - cuda_device::barrier::barrier_cta_sync(1, 128); + async_copy::cp_async_wait_group::<0>(); + mbarrier::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(); @@ -285,8 +335,8 @@ impl TilePipeline { b = next_b; }); scores.mma(a, b); - async_copy::cp_async_wait_group(0); - cuda_device::barrier::barrier_cta_sync(1, 128); + async_copy::cp_async_wait_group::<0>(); + mbarrier::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(); @@ -304,10 +354,12 @@ impl TilePipeline { } } +/// 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))] @@ -389,7 +441,7 @@ pub mod kernels { pipeline.step::(remaining, &mut output, &mut state); } output.store(&state, &pipeline.q, warp_id, lane as usize); - cuda_device::barrier::barrier_cta_sync(1, 128); + mbarrier::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; diff --git a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/main.rs b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/main.rs index 2b9a62a5a8..125b6c30d0 100644 --- a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/main.rs +++ b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/main.rs @@ -3,13 +3,18 @@ * 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/cuteir/examples/fmha_prefill_sm120a_cute/src/runner.rs b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/runner.rs index 79cb8c4304..e8911fcc85 100644 --- a/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/runner.rs +++ b/cuda-oxide/cuteir/examples/fmha_prefill_sm120a_cute/src/runner.rs @@ -3,6 +3,8 @@ * 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}; @@ -10,36 +12,57 @@ 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, @@ -129,6 +152,7 @@ impl Args { } } +/// Dispatch the comparison process or a standalone prefill run. pub fn run() -> Result<()> { let args = Args::parse()?; if !args.rust_only { @@ -158,6 +182,7 @@ pub fn run() -> Result<()> { } } +/// 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), @@ -165,6 +190,7 @@ fn run_mask(args: Args) -> Result<()> { } } +/// Verify the kernel and replay, collect CUDA-event samples, and write the report. fn run_shape( args: Args, ) -> Result<()> { @@ -298,14 +324,17 @@ fn run_shape 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()); @@ -314,12 +343,14 @@ mod tests { 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"] { @@ -329,6 +360,7 @@ mod tests { } } + /// Reject standalone-only options in a multi-case comparison. #[test] fn all_cases_rejects_standalone_execution_and_output() { for options in [ @@ -340,6 +372,7 @@ mod tests { } } + /// 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")] { @@ -347,6 +380,7 @@ mod tests { } } + /// 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"] { @@ -355,6 +389,7 @@ mod tests { } } + /// 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")] { diff --git a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_ctm/src/kernel.rs b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_ctm/src/kernel.rs index f8962c16e3..6657c2e42b 100644 --- a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_ctm/src/kernel.rs +++ b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_ctm/src/kernel.rs @@ -26,34 +26,59 @@ use cuda_device::{ DynamicSharedArray, cluster_launch, cuda_module, kernel, launch_bounds, launch_contract, }; +/// Per-CTA 128x64 FP16 A tile with the shared swizzle used by its tensor map. pub type ASmem = F16Tile128; +/// Per-CTA 96x64 FP16 B tile feeding the 192-column MMA partition. pub type B0Smem = F16TileB96; +/// Per-CTA 80x64 FP16 B tile feeding the 160-column MMA partition. pub type B1Smem = F16TileB80; +/// Per-CTA 128x32 FP16 output tile with the B64 shared swizzle. pub type CSmem = F16TileC; +/// Two-CTA FP16 MMA shape for the 256x192 left output partition. type Left = DenseSharedF16; +/// Two-CTA FP16 MMA shape for the 256x160 right output partition. type Right = DenseSharedF16; +/// Ready two-CTA tensor-memory allocation retaining all 512 allocated columns. type Accumulator = TmemGuard; +/// Borrowed FP32 tensor-memory view of the 256x192 left accumulator partition. type LeftView<'a> = TmemView<'a, TmemF32<256, 192>, 512, CtaGroup2>; +/// Borrowed FP32 tensor-memory view of the 256x160 right accumulator partition. type RightView<'a> = TmemView<'a, TmemF32<256, 160>, 512, CtaGroup2>; +/// Eight FP16 values stored as one naturally aligned 16-byte shared-memory chunk. #[repr(C, align(16))] -struct F16x8([f16; 8]); +struct F16x8( + #[doc = "FP16 lanes in increasing column order within the contiguous chunk."] [f16; 8], +); +/// Compile-time check that the chunk carrier has exactly 16-byte size and alignment. const _: () = assert!(core::mem::size_of::() == 16 && core::mem::align_of::() == 16); +/// 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 = ASmem::FOOTPRINT_BYTES; +/// Shared bytes occupied by both B partitions in one CTA stage. const B_BYTES: usize = B0Smem::FOOTPRINT_BYTES + B1Smem::FOOTPRINT_BYTES; +/// 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 = CSmem::FOOTPRINT_BYTES; +/// 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; +/// Combined A/B input bytes delivered by both CTAs for one K64 stage. const TX_BYTES: u32 = (2 * (A_BYTES + B_BYTES)) as u32; +/// Compile-time checks for the complete shared-memory footprint and two-CTA transaction bytes. const _: () = assert!(DYNAMIC_SMEM_BYTES == 211_072 && TX_BYTES == 77_824); // A violated physical contract must abort the kernel, never let one lane leave // a collective while its peers wait. All valid launch layouts satisfy these. +/// Trap on a physical-contract error so one lane cannot leave its collective peers waiting. #[inline(always)] fn checked(value: Result) -> T { match value { @@ -62,25 +87,30 @@ fn checked(value: Result) -> T { } } +/// Borrow a typed shared tile; the caller provides live, correctly aligned storage of at least its footprint. #[inline(always)] unsafe fn region<'s, F: PhysicalFormat>(base: *mut u8) -> SharedRegion<'s, F> { unsafe { SharedRegion::::from_raw_parts_or_trap(base.cast(), F::FOOTPRINT_BYTES) } } +/// Initialize a live barrier with the requested arrivals; the caller owns initialization and publication ordering. #[inline(always)] unsafe fn init(pointer: *mut Barrier, arrivals: u32) { unsafe { ManagedBarrier::::from_static(pointer).init_instruction(arrivals); } } +/// Retry cluster parity probes on a live initialized barrier until the requested phase completes. #[inline(always)] unsafe fn wait(barrier: &ManagedBarrier, phase: u32) { unsafe { while !mbarrier::try_wait_parity_cluster(barrier, phase) {} } } +/// Return the checked two-bit participation mask naming both CTAs in the cluster. #[inline(always)] fn both_ctas() -> CtaMask { checked(CtaMask::try_from_bits(3)) } +/// Issue one K16 step for both N partitions with fill/last-use A collectors; operands must be ready and all issuer lanes participate. #[inline(always)] #[allow(clippy::too_many_arguments)] unsafe fn issue_pair( @@ -108,6 +138,7 @@ unsafe fn issue_pair( } } +/// Issue all four K16 steps for a ready shared stage; stage must index the live five-slot allocation. #[inline(always)] unsafe fn issue_stage( base: *mut u8, @@ -133,6 +164,7 @@ unsafe fn issue_stage( } // Each swizzled 8-half chunk is contiguous and 16-byte aligned. +/// Add bias and store eight FP16 lanes; output must be 16-byte aligned, tid below 128, and COL a valid eight-lane chunk. #[inline(always)] unsafe fn store_chunk(output: *mut f16, values: [f32; 32], tid: u32, bias: f32) { unsafe { @@ -152,6 +184,7 @@ unsafe fn store_chunk(output: *mut f16, values: [f32; 32], tid } } +/// Read a ready TMEM subtile, wait for its load, and store biased FP16 chunks into live aligned shared output. #[inline(always)] unsafe fn epilogue(acc: &Accumulator, output: *mut f16, tid: u32, subtile: u32, bias: f32) { unsafe { @@ -167,8 +200,10 @@ unsafe fn epilogue(acc: &Accumulator, output: *mut f16, tid: u32, subtile: u32, } } +/// Two-CTA SM100 GEMM entry point and its explicit launch/storage contract. #[cuda_module] pub mod kernels { + use super::*; /// Tensor allocations and all four map descriptors stay live through completion. diff --git a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_ctm/src/main.rs b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_ctm/src/main.rs index 2de7bd545a..75fe701f21 100644 --- a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_ctm/src/main.rs +++ b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_ctm/src/main.rs @@ -3,6 +3,8 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Host input preparation, validation, and timing for the two-CTA 256x352 FP16 GEMM. + #![feature(f16)] mod kernel; @@ -17,15 +19,20 @@ use std::{ 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", @@ -33,6 +40,7 @@ impl TimingMode { } } + /// Report the CUDA timing mechanism used for the selected mode. fn tool(self) -> &'static str { match self { Self::Direct => "cuda-events-single-kernel", @@ -41,23 +49,38 @@ impl TimingMode { } } +/// Validated matrix shape, timing controls, and input/output paths for the benchmark. struct Args { + /// 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 { m: 256, @@ -175,6 +198,7 @@ impl Args { } } +/// 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(()) @@ -185,12 +209,16 @@ fn cuda_status(status: sys::CUresult, operation: &str) -> Result<()> { // 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, @@ -246,6 +274,7 @@ impl KernelGraph { }) } + /// Bind the stream's context and enqueue the captured graph once. fn replay(&self) -> Result<()> { self.stream.context().bind_to_thread()?; cuda_status( @@ -265,6 +294,7 @@ impl Drop for KernelGraph { } } +/// 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")? { @@ -283,6 +313,7 @@ fn load_f16(path: &Path, count: usize) -> Result> { .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)?); @@ -295,6 +326,7 @@ fn store_f16(path: &Path, data: &[u16]) -> Result<()> { // The standalone run uses reproducible integers in [-2, 2), exact in FP16. // External-oracle runs may supply shared input bytes to multiple 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(|_| { @@ -307,6 +339,7 @@ fn integer_input(count: usize, mut state: u32) -> Vec { .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; @@ -339,6 +372,7 @@ fn verify(args: &Args, a: &[u16], b: &[u16], bias: &[u16], got: &[u16]) -> Resul Ok(()) } +/// Validate inputs, launch the SM100 kernel, verify results, and report direct or graph timing. fn main() -> Result<()> { let args = Args::parse()?; let single_launch = diff --git a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_ctm/src/physical.rs b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_ctm/src/physical.rs index 6febc2a622..511a85f54c 100644 --- a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_ctm/src/physical.rs +++ b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_ctm/src/physical.rs @@ -1,8 +1,13 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +//! FP16 output swizzle shared by the device store path and exhaustive CPU layout checks. + #[cfg(test)] mod tests { + //! Enumerate output positions to check complete coverage and contiguous vector-store chunks. + use super::*; + /// Check every output element appears once and every eight-lane chunk stays contiguous under the swizzle. #[test] fn output_swizzle_covers_each_element_once_and_preserves_eight_element_chunks() { let mut seen = [false; 128 * 32]; diff --git a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/README.md b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/README.md index d29f11cf61..750f827f50 100644 --- a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/README.md +++ b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/README.md @@ -214,7 +214,7 @@ 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. -- [SM100 APIs](../../cute-rs/src/sm100.rs): `Sm100SharedTile`, +- [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/sm100_ops.rs): layout, collector, diff --git a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/src/kernel.rs b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/src/kernel.rs index 7cb7d401ce..37b96c4895 100644 --- a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/src/kernel.rs +++ b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/src/kernel.rs @@ -30,28 +30,46 @@ use cute_rs::{ // 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 = Composed, 0, RowMajor<128, 64>>; +/// Per-CTA 96x64 FP16 B tile feeding the 192-column MMA partition. pub type B0Smem = Composed, 0, RowMajor<96, 64>>; +/// Per-CTA 80x64 FP16 B tile feeding the 160-column MMA partition. pub type B1Smem = Composed, 0, RowMajor<80, 64>>; +/// Per-CTA 128x32 FP16 output tile with the B64 shared swizzle. pub type CSmem = Composed, 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 = Sm100TmaMmaPipeline<5, 77_824>; +/// Accumulator handoff between the MMA producer and four epilogue warps per CTA. type AccumulatorPipeline = Sm100AccumulatorPipeline<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 { @@ -60,8 +78,10 @@ unsafe fn wait_deallocation(bar: *const Barrier) { } } +/// 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. diff --git a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/src/main.rs b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/src/main.rs index f6d16e2bd5..eeb3679b7c 100644 --- a/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/src/main.rs +++ b/cuda-oxide/cuteir/examples/fp16_gemm_256x352_cute/src/main.rs @@ -3,6 +3,8 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Host input preparation, validation, and timing for the two-CTA 256x352 FP16 GEMM. + #![feature(f16)] mod kernel; @@ -16,15 +18,20 @@ use std::{ 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", @@ -32,6 +39,7 @@ impl TimingMode { } } + /// Report the CUDA timing mechanism used for the selected mode. fn tool(self) -> &'static str { match self { Self::Direct => "cuda-events-single-kernel", @@ -40,25 +48,42 @@ impl TimingMode { } } +/// 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, @@ -199,6 +224,7 @@ impl 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"); @@ -253,6 +279,7 @@ fn compare(args: &Args) -> Result<()> { 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(()) @@ -263,12 +290,16 @@ fn cuda_status(status: sys::CUresult, operation: &str) -> Result<()> { // 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, @@ -324,6 +355,7 @@ impl KernelGraph { }) } + /// Bind the stream's context and enqueue the captured graph once. fn replay(&self) -> Result<()> { self.stream.context().bind_to_thread()?; cuda_status( @@ -343,6 +375,7 @@ impl Drop for KernelGraph { } } +/// 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")? { @@ -361,6 +394,7 @@ fn load_f16(path: &Path, count: usize) -> Result> { .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)?); @@ -373,6 +407,7 @@ fn store_f16(path: &Path, data: &[u16]) -> Result<()> { // 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(|_| { @@ -385,6 +420,7 @@ fn integer_input(count: usize, mut state: u32) -> Vec { .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; @@ -417,6 +453,7 @@ fn verify(args: &Args, a: &[u16], b: &[u16], bias: &[u16], got: &[u16]) -> Resul 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 { diff --git a/cuda-oxide/cuteir/examples/nvfp4_gemv_cute/src/main.rs b/cuda-oxide/cuteir/examples/nvfp4_gemv_cute/src/main.rs index f22de59bdf..86a004644b 100644 --- a/cuda-oxide/cuteir/examples/nvfp4_gemv_cute/src/main.rs +++ b/cuda-oxide/cuteir/examples/nvfp4_gemv_cute/src/main.rs @@ -24,17 +24,26 @@ use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; use cute_rs::{BlockScaledTensor, E2M1, KMajor, Mkl, Nkl, Sm1xxBlockScaleKMajor, UE8M0}; +/// 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, E2M1, UE8M0, SF_VEC, KMajor>; +/// Packed E2M1 vector with UE8M0 scales and K-major vector coordinates. type VectorB<'a> = BlockScaledTensor<'a, E2M1, UE8M0, SF_VEC, KMajor>; +/// Device entry point for one-thread-per-row block-scaled GEMV. #[cuda_module] mod kernels { use super::*; @@ -95,19 +104,26 @@ mod kernels { } } +/// 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, @@ -144,10 +160,12 @@ fn parse_args() -> Args { 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() @@ -155,6 +173,7 @@ fn checked_product(label: &str, factors: &[usize]) -> usize { .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]); @@ -165,7 +184,9 @@ fn validate_problem_sizes(m: usize, k: usize, l: usize) { 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 { @@ -175,6 +196,7 @@ fn e2m1_to_f32(nibble: u8) -> f32 { } } +/// 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) @@ -185,22 +207,27 @@ fn ue8m0_to_f32(byte: u8) -> f32 { } } +/// 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 { @@ -215,6 +242,7 @@ fn pack_a(m: usize, k: usize, l: usize) -> Vec { 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 { @@ -227,14 +255,17 @@ fn pack_b(k: usize, l: usize) -> Vec { 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).offset(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. @@ -253,6 +284,7 @@ fn make_scales(m: usize, k: usize, l: usize) -> (Vec, Vec) { (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 { @@ -271,6 +303,7 @@ fn reference(m: usize, k: usize, l: usize) -> Vec { 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() { @@ -291,12 +324,14 @@ fn validate(got: &[f16], expected: &[f16]) { ); } +/// 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!( @@ -361,11 +396,13 @@ fn main() { 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]; @@ -377,6 +414,7 @@ mod tests { 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); @@ -388,6 +426,7 @@ mod tests { 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); @@ -416,6 +455,7 @@ mod tests { 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); diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/src/kernel.rs b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/src/kernel.rs index 62a4525bdd..58b479bc7b 100644 --- a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/src/kernel.rs +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/src/kernel.rs @@ -19,20 +19,34 @@ use cuda_device::{ use cute_rs::sm100_attention::{MmaSs, MmaTs, SharedLayout, SharedTile, TmaCopy, TmemFragment}; use cute_rs::{OperandA, OperandB, TmaDesc}; +/// K-major shared layout for a 128 × 128 FP16 query tile. pub type QLayout = SharedLayout<128, 128>; +/// K-major shared layout for a 128 × 128 FP16 key tile. pub type KLayout = SharedLayout<128, 128>; +/// MN-major shared layout for a 128 × 128 FP16 value tile. pub type VLayout = SharedLayout<128, 128, true>; +/// K-major shared layout used to store a 128 × 128 FP16 output tile. pub type OLayout = SharedLayout<128, 128>; +/// Single-CTA 128 × 128 × 128 QK MMA with both operands in shared memory. type Qk = MmaSs; +/// Single-CTA 128 × 128 × 128 PV MMA with probabilities in tensor memory. type Pv = MmaTs; +/// Threads per CTA, partitioned among softmax, correction, MMA, and TMA roles. pub const THREADS: u32 = 512; +/// Bytes in one 128 × 128 FP16 shared-memory tile. const TILE_BYTES: usize = 128 * 128 * 2; +/// Byte offset of the two output tiles, after the barrier/token prefix. const O_OFFSET: usize = 1024; +/// Byte offset of the two query tiles in dynamic shared memory. const Q_OFFSET: usize = O_OFFSET + 2 * TILE_BYTES; +/// Byte offset of the three shared slots that alternate key and value tiles. const KV_OFFSET: usize = Q_OFFSET + 2 * TILE_BYTES; +/// Total dynamic shared-memory bytes required by the launch contract. pub const DYNAMIC_SMEM_BYTES: usize = KV_OFFSET + 3 * TILE_BYTES; +/// Byte offset of the TMEM allocation token, after 34 eight-byte barriers. const TMEM_TOKEN_OFFSET: usize = 34 * 8; +/// Attention scale log2(e) / sqrt(128), allowing softmax to use base-two exponentials. const SCALE_LOG2: f32 = 0.127_517_43; // log2(e) / sqrt(128) /// One independent full/empty stage ring. A producer count starts at phase1 @@ -40,11 +54,15 @@ const SCALE_LOG2: f32 = 0.127_517_43; // log2(e) / sqrt(128) /// The explicit arrival counts below are the source's participant contracts. #[derive(Clone, Copy)] struct Ring { + /// Contiguous stage barriers completed when producers publish their work. full: *mut Barrier, + /// Contiguous stage barriers completed when consumers release their slots. empty: *mut Barrier, } impl Ring { + /// Bind a ring at `offset` measured in eight-byte barrier slots. + /// The shared allocation must contain `2 * STAGES` aligned, nonoverlapping slots there. #[inline(always)] unsafe fn new(storage: *mut u8, offset: usize) -> Self { let full = unsafe { storage.add(offset * 8).cast::() }; @@ -53,6 +71,8 @@ impl Ring { empty: unsafe { full.add(STAGES as usize) }, } } + /// Initialize full/empty barriers with per-stage producer/consumer arrival counts. + /// Call once before publishing initialization to the CTA. #[inline(always)] unsafe fn init(self, producers: u32, consumers: u32) { let mut i = 0; @@ -64,20 +84,24 @@ impl Ring { i += 1; } } + /// Wait for consumers to release `count % STAGES`, using the producer’s inverted parity. #[inline(always)] unsafe fn acquire(self, count: u32) { let ptr = unsafe { self.empty.add((count % STAGES) as usize) }; unsafe { wait(ptr, ((count / STAGES) & 1) ^ 1) }; } + /// Wait for producers to fill `count % STAGES`, then order subsequent TMEM accesses. #[inline(always)] unsafe fn wait(self, count: u32) { unsafe { wait(self.full_ptr(count), (count / STAGES) & 1) }; tcgen05::tcgen05_fence_after_thread_sync(); } + /// Select the full barrier for the monotonically increasing stage count. #[inline(always)] unsafe fn full_ptr(self, count: u32) -> *mut Barrier { unsafe { self.full.add((count % STAGES) as usize) } } + /// Order prior TMEM accesses and contribute one thread arrival to the full barrier. #[inline(always)] unsafe fn commit_threads(self, count: u32) { tcgen05::tcgen05_fence_before_thread_sync(); @@ -85,6 +109,7 @@ impl Ring { barrier::mbarrier_arrive(self.full_ptr(count)); } } + /// Order prior TMEM accesses and contribute one thread arrival to the empty barrier. #[inline(always)] unsafe fn release_threads(self, count: u32) { tcgen05::tcgen05_fence_before_thread_sync(); @@ -92,6 +117,8 @@ impl Ring { barrier::mbarrier_arrive(self.empty.add((count % STAGES) as usize)); } } + /// Have one elected lane signal the full barrier after outstanding MMA work completes. + /// All lanes of the MMA warp must participate in the election. #[inline(always)] unsafe fn commit_mma(self, count: u32) { if warp::is_elected_sync(u32::MAX) { @@ -100,6 +127,8 @@ impl Ring { } } } + /// Have one elected lane release the slot after outstanding MMA work completes. + /// All lanes of the MMA warp must participate in the election. #[inline(always)] unsafe fn release_mma(self, count: u32) { if warp::is_elected_sync(u32::MAX) { @@ -108,6 +137,8 @@ impl Ring { } } } + /// Have one elected lane arrive and expect one FP16 tile’s TMA byte count. + /// All lanes of the producer warp must participate in the election. #[inline(always)] unsafe fn expect(self, count: u32) { if warp::is_elected_sync(u32::MAX) { @@ -118,6 +149,7 @@ impl Ring { } } +/// Spin until an initialized shared barrier completes the requested parity phase. #[inline(always)] unsafe fn wait(ptr: *const Barrier, phase: u32) { while !unsafe { barrier::mbarrier_try_wait_parity(ptr, phase) } {} @@ -126,21 +158,34 @@ unsafe fn wait(ptr: *const Barrier, phase: u32) { /// A copy of the same twelve pipeline identities exists in every participant. #[derive(Clone, Copy)] struct Pipelines { + /// Two query stages transferred by TMA and released by the MMA warp. q: Ring<2>, + /// Three alternating key/value stages transferred by TMA and released by MMA. kv: Ring<3>, + /// QK scores from MMA to the first 128-thread softmax group. s0: Ring<1>, + /// QK scores from MMA to the second 128-thread softmax group. s1: Ring<1>, + /// FP16 probabilities from the first softmax group to PV MMA. p0: Ring<1>, + /// FP16 probabilities from the second softmax group to PV MMA. p1: Ring<1>, + /// First-group statistics rendezvous: 32 MMA plus 128 softmax arrivals, then correction. stats0: Ring<1>, + /// Second-group statistics rendezvous: 32 MMA plus 128 softmax arrivals, then correction. stats1: Ring<1>, + /// Two corrected output stages handed from 128 correction threads to 32 TMA-store threads. epi: Ring<2>, + /// Two PV accumulator stages handed from MMA completion to 128 correction threads. output: Ring<2>, + /// First softmax group’s handoff protecting aliased score/probability TMEM from the second. alias0: Ring<1>, + /// Second softmax group’s handoff protecting aliased score/probability TMEM from the first. alias1: Ring<1>, } impl Pipelines { + /// Bind all twelve rings to the first 34 aligned barrier slots in live CTA shared memory. #[inline(always)] unsafe fn new(smem: *mut u8) -> Self { unsafe { @@ -160,6 +205,7 @@ impl Pipelines { } } } + /// Initialize the schedule’s producer/consumer counts once before CTA-wide publication. #[inline(always)] unsafe fn init(self) { unsafe { @@ -180,6 +226,7 @@ impl Pipelines { } } +/// Count 128-key tiles for the 256-query tile starting at `row`, clipping causal work. #[inline(always)] fn trip_count(row: u32, keys: u32, causal: bool) -> u32 { let limit = if causal && row + 256 < keys { @@ -190,11 +237,13 @@ fn trip_count(row: u32, keys: u32, causal: bool) -> u32 { limit.div_ceil(128) } +/// Add the calling warp’s 32-row offset within a 128-thread group to a TMEM address. #[inline(always)] fn tmem_rows(base: u32, tid: u32) -> u32 { base + ((((tid % 128) / 32) * 32) << 16) } +/// Evaluate the device approximate base-two exponential with subnormal flushing. #[inline(always)] fn exp2(x: f32) -> f32 { cuda_device::float::ex2_approx_ftz_f32(x) @@ -234,6 +283,7 @@ fn reduce_masked_scores(scores: &TmemFragment<128>) -> f32 { maxima[0] } +/// Reduce the first `2 * WIDTH` maxima in place to `WIDTH` pairwise maxima; `WIDTH <= 32`. #[inline(always)] fn reduce_max_level(maxima: &mut [f32; 64]) { let mut i = 0; @@ -244,12 +294,15 @@ fn reduce_max_level(maxima: &mut [f32; 64]) { } } +/// Multiply two FP32 values by packed FP32 factors and return their packed FP16 bits. #[inline(always)] fn scaled_half_pair(lo: f32, hi: f32, factor: u64) -> u32 { let scaled = mul_f32x2(pack_f32x2(lo, hi), factor); cuda_device::convert::cvt_f16x2_f32(unpack_f32x2_lo(scaled), unpack_f32x2_hi(scaled)) } +/// Acquire score/statistics slots and issue QK without accumulation, then publish score completion. +/// All 32 MMA-warp lanes participate; `destination` is a live TMEM address. #[inline(always)] unsafe fn mma_qk( destination: u32, @@ -271,6 +324,8 @@ unsafe fn mma_qk( } } +/// Wait for probabilities, issue PV into an acquired output slot, and release P after MMA completion. +/// All 32 MMA-warp lanes participate; TMEM operands must remain live through completion. #[inline(always)] #[allow(clippy::too_many_arguments)] // Keep the source pipeline handoffs explicit. unsafe fn mma_pv( @@ -294,6 +349,7 @@ unsafe fn mma_pv( } } +/// Device entry point and generated host launch bindings for the persistent SM100 schedule. #[cuda_module] pub mod kernels { use super::*; diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/src/main.rs b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/src/main.rs index 38423d7e27..f2473caefa 100644 --- a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/src/main.rs +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/src/main.rs @@ -2,6 +2,10 @@ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ + +//! SM100 FP16 prefill comparison runner and file-based CUDA adapter. +//! Adapter mode validates against an FP32 oracle before collecting CUDA-event timings. + #![feature(f16)] mod kernel; @@ -11,23 +15,37 @@ use cute_rs::sm100_attention::{SmemLayout, make_tma_desc}; use serde_json::json; use std::{error::Error, path::PathBuf}; +/// Fallible host operation carrying CLI, file, Python, or CUDA errors. type Result = std::result::Result>; +/// Validated comparison options and optional file-based adapter inputs. struct Args { + /// Batch, query length, key length, query heads, and KV heads; head dimension is 128. shape: [u32; 5], + /// Enable top-left causal masking instead of dense attention. causal: bool, + /// Untimed launches before each timing window. warmup: usize, + /// Timed launches per CUDA-event window; must be nonzero. iters: usize, + /// Number of independent timing windows; must be nonzero. windows: usize, + /// Validate output without collecting timing samples. check_only: bool, + /// Adapter directory containing little-endian q.f16, k.f16, and v.f16 files. input_dir: Option, + /// Adapter output path for the validated little-endian FP16 tensor. output: Option, + /// Adapter input path for the little-endian FP32 reference tensor. oracle: Option, + /// Optional destination for the comparison or adapter report. json: Option, + /// Absolute error tolerance in [0, 0.1], combined with a fixed 1e-5 relative tolerance. atol: f32, } impl Args { + /// Parse CLI options and reject incompatible heads, overflowing shapes, or invalid timing/tolerance values. fn parse() -> Result { let mut out = Self { shape: [1, 1024, 1024, 8, 8], @@ -120,6 +138,7 @@ impl Args { } } +/// Run the pinned Python comparison harness, forwarding shape and timing options. fn compare(args: &Args) -> Result<()> { let directory = PathBuf::from(env!("CARGO_MANIFEST_DIR")); let python = std::env::var_os("CUDA_OXIDE_CUTE_PYTHON") @@ -170,6 +189,7 @@ fn compare(args: &Args) -> Result<()> { Ok(()) } +/// Read exactly `count` little-endian FP16 elements, preserving their raw bits. fn load_half(path: &std::path::Path, count: usize) -> Result> { let data = std::fs::read(path)?; if data.len() != count.checked_mul(2).ok_or("FP16 byte count overflow")? { @@ -183,6 +203,7 @@ fn load_half(path: &std::path::Path, count: usize) -> Result> { .collect()) } +/// Pack one 128-byte TMA descriptor per batch/head for a BSHD allocation with D=128. fn descriptors(base: u64, batch: u32, sequence: u32, heads: u32) -> Result> { let mut result = Vec::with_capacity(batch as usize * heads as usize * 128); for b in 0..batch { @@ -201,6 +222,8 @@ fn descriptors(base: u64, batch: u32, sequence: u32, heads: u32) Ok(result) } +/// Check finite FP16 outputs against FP32 values using `atol + 1e-5 * abs(expected)`. +/// Return maximum absolute error and root-mean-square error, or reject any mismatch. fn verify(output: &[u16], oracle: &[f32], atol: f32) -> Result<(f32, f64)> { if output.len() != oracle.len() { return Err("oracle and output element counts differ".into()); @@ -235,6 +258,7 @@ fn verify(output: &[u16], oracle: &[f32], atol: f32) -> Result<(f32, f64)> { Ok((maximum, (square / output.len() as f64).sqrt())) } +/// Run the comparison harness, or validate, time, and save the SM100 adapter’s output. fn main() -> Result<()> { let args = Args::parse()?; let Some(directory) = &args.input_dir else { diff --git a/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/attention-feature-gaps.md b/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/attention-feature-gaps.md index 9d7ece4887..f9ca477b0f 100644 --- a/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/attention-feature-gaps.md +++ b/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/attention-feature-gaps.md @@ -48,15 +48,15 @@ for the implemented configurations and their recorded results. ## Implementation entry points -- [Attention types and layout constraints](../../../cute-rs/src/sm100_attention.rs), +- [Attention types and layout constraints](../../../cute-rs/src/sm100_attention/mod.rs), [host TMA descriptors](../../../cute-rs/src/sm100_attention/host.rs), and [CuTe operation contracts](../../../dialect-cute/src/sm100_attention_ops.rs) currently specialize the attention path to FP16. -- [CuTe importing](../../../../crates/mir-importer/src/translator/cute/sm100_attention_emit.rs) - and [attention export](../../../../crates/cuda-oxide-mlir-export/src/cute_sm100_attention.rs) +- [CuTe importing](../../../../crates/mir-importer/src/translator/cute/sm100_attention.rs) + and [attention export](../../../../crates/cuda-oxide-mlir-export/src/cute/independent_sm100.rs) must preserve any new element and layout information through translation. -- [Raw MMA export](../../../../crates/cuda-oxide-mlir-export/src/nvvm_tcgen05.rs) - and [numeric export](../../../../crates/cuda-oxide-mlir-export/src/nvvm_numeric.rs) +- [Raw MMA export](../../../../crates/cuda-oxide-mlir-export/src/nvvm/tcgen05.rs) + and [numeric export](../../../../crates/cuda-oxide-mlir-export/src/nvvm/numeric.rs) are the entry points for additional MMA kinds, conversions, and packed arithmetic. Review the [selected-backend registry](../../../../crates/cuda-oxide-mlir-export/src/profile.rs) when adding an operation that already has a native NVVM route. diff --git a/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/measurement-notes.md b/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/measurement-notes.md index 712fd76b37..916729500c 100644 --- a/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/measurement-notes.md +++ b/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/measurement-notes.md @@ -87,7 +87,7 @@ The example still expresses MMA using layout- and operand-typed CuTe tiles and TMA using typed tensor maps. It contains no hand-written PTX kernel. New reusable abstractions live in CuTe-RS: -- [`SharedTileArray`](../../../cute-rs/src/sm100_attention.rs) carries layout, operand +- [SharedTileArray](../../../cute-rs/src/sm100_attention/shared.rs) carries layout, operand role and extent while exposing complete, swizzle-preserving tile views. - [`TwoCtaRing` and `RingCursor`](../../../cute-rs/src/sm100_attention/pipeline.rs) encode stage count, cursor role and permitted synchronization/completion policies. diff --git a/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/src/comparison.rs b/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/src/comparison.rs index a347154a6e..c2f8df4226 100644 --- a/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/src/comparison.rs +++ b/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/src/comparison.rs @@ -3,12 +3,15 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! Default CuTeDSL comparison dispatch and the serialized-input adapter boundary. + use std::{ ffi::OsString, path::{Path, PathBuf}, process::Command, }; +/// Run help or the Python comparison; return false only for serialized-input adapter mode. pub fn dispatch(arguments: &[OsString]) -> super::Result { // compare.py re-enters this executable with serialized inputs for the Rust run. if arguments.iter().any(|argument| argument == "--input-dir") { @@ -40,6 +43,7 @@ pub fn dispatch(arguments: &[OsString]) -> super::Result { Ok(true) } +/// Check local reference assets and forward arguments, pinning the currently running Rust binary. fn comparison_command( directory: &Path, python: &Path, @@ -64,6 +68,7 @@ fn comparison_command( Ok(command) } +/// Run the comparison subprocess and report startup or nonzero-exit errors. fn run(command: &mut Command) -> super::Result<()> { let status = command.status().map_err(|error| { format!( @@ -77,17 +82,22 @@ fn run(command: &mut Command) -> super::Result<()> { Ok(()) } +/// Check argument forwarding and adapter dispatch without launching Python or CUDA. #[cfg(test)] mod tests { use super::*; use std::{fs, os::unix::fs::PermissionsExt}; + /// Temporary executable that records arguments and simulates a comparison exit status. struct FakePython { + /// Private temporary directory holding reference placeholders, the executable and captured arguments. directory: PathBuf, } impl FakePython { + /// Create an isolated recording executable and minimal comparison assets. fn new() -> Self { + /// Per-process sequence preventing temporary-directory collisions between concurrent tests. static NEXT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); let directory = std::env::temp_dir().join(format!( "decode-cli-{}-{}", @@ -107,6 +117,7 @@ mod tests { Self { directory } } + /// Build a comparison command whose simulated interpreter records every argument. fn command(&self, arguments: &[OsString]) -> Command { let mut command = comparison_command( &self.directory, @@ -126,6 +137,7 @@ mod tests { } } + /// Forward every option unchanged and keep a binary path containing spaces as one argument. #[test] fn comparison_forwards_options_and_preserves_binary_paths_with_spaces() { let python = FakePython::new(); @@ -171,6 +183,7 @@ mod tests { ); } + /// Propagate a nonzero comparison exit while preserving evidence that it was launched. #[test] fn failed_comparison_is_an_error() { let python = FakePython::new(); @@ -180,6 +193,7 @@ mod tests { assert!(python.directory.join("arguments").is_file()); } + /// Serialized input arguments must bypass the Python comparison dispatcher. #[test] fn adapter_input_does_not_reenter_comparison() { let arguments = ["--input-dir", "/inputs", "--output-dir", "/outputs"].map(OsString::from); diff --git a/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/src/kernel.rs b/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/src/kernel.rs index 5df2ccd87b..9c57e935a3 100644 --- a/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/src/kernel.rs +++ b/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/src/kernel.rs @@ -29,35 +29,61 @@ use cute_rs::sm100_attention::{ }; use cute_rs::{OperandA, OperandB, TmaDesc}; +/// Per-CTA 64-by-64 FP16 query tile, stored K-major for QK MMA. pub type QLayout = SharedLayout<64, 64>; +/// Per-CTA 64-by-64 FP16 cache-key tile, stored K-major for QK MMA. pub type KLayout = SharedLayout<64, 64>; +/// Per-CTA 64-by-32 FP16 probability chunk consumed by PV MMA. pub type PLayout = SharedLayout<64, 32>; +/// Per-CTA 128-by-32 MN-major FP16 value chunk consumed by PV MMA. pub type VLayout = SharedLayout<128, 32, true>; +/// Group-two M128/N128/K64 QK instruction using shared query and key tiles. type QkMma = MmaSs; +/// Group-two M128/N256/K32 PV instruction using shared probability and value tiles. type PvMma = MmaSs; +/// Threads in each persistent CTA, including dedicated load, MMA and reduction warps. pub const THREADS: u32 = 384; +/// Required dynamic shared-memory bytes per CTA, with a 1024-byte-aligned base. pub const DYNAMIC_SMEM_BYTES: u32 = 232_448; +/// Shared byte offset of eight latent-query chunks. const Q: usize = 0; +/// Shared byte offset of the rotary-query chunk following the latent query. const QR: usize = 65_536; +/// Shared byte offset of fifteen reusable key/value stages, each 8192 bytes. const KV: usize = 73_728; +/// Shared byte offset of two probability buffers, each containing four K32 chunks. const P: usize = 196_608; +/// Shared byte offset of four page-index slots, spaced 256 bytes apart. const PT: usize = 229_376; +/// Shared byte offset of two 128-element FP32 softmax-maximum exchange buffers. const MAX_EXCHANGE: usize = 230_400; +/// Shared byte offset of 128 FP32 partial sums used during final normalization. const SUM_EXCHANGE: usize = 231_424; +/// Shared byte offset of the 54 full/empty mbarriers backing all pipeline rings. const BARS: usize = 231_936; +/// Shared byte offset of the u32 slot receiving this CTA’s TMEM allocation address. const TOKEN: usize = BARS + 54 * 8; +/// Convert the natural-exponential softmax scale to the base-2 instruction scale. const LOG2_E: f32 = core::f32::consts::LOG2_E; // Each pipeline carries its stage count and full/empty memory protocol in its type. +/// Single query buffer: TMA fills it, then MMA completion releases it for reuse. type QueryRing = TwoCtaRing<1, AsyncCompletion, AsyncCopyReuse>; +/// Fifteen shared stages reused by TMA producers after MMA finishes reading each tile. type KvRing = TwoCtaRing<15, AsyncCompletion, AsyncCopyReuse>; +/// Two TMEM score buffers passed from QK MMA to softmax threads. type ScoreRing = TwoCtaRing<2, TensorMemory>; +/// Two shared probability buffers published by softmax and released after PV MMA. type ProbabilityRing = TwoCtaRing<2, SharedToAsync, SharedReuse>; +/// Two local metadata buffers passed from softmax to output-correction threads. type CorrectionRing = TwoCtaRing<2, LocalMemory>; +/// Single TMEM output accumulator handed between PV MMA and correction threads. type OutputRing = TwoCtaRing<1, TensorMemory>; +/// Four local shared page-index slots filled by asynchronous copies and read by the load warp. type PageRing = TwoCtaRing<4, LocalShared>; +/// Select the larger score or log-sum-exp value with Rust FP32 max semantics. #[inline(always)] fn maxf(a: f32, b: f32) -> f32 { a.max(b) @@ -66,6 +92,7 @@ fn maxf(a: f32, b: f32) -> f32 { /// A ternary reduction exposes independent maximum operations. #[inline(always)] fn score_max_ternary(values: &[f32; 64], previous: f32) -> f32 { + /// Reduce three scores while exposing a balanced ternary maximum tree. #[inline(always)] fn max3(a: f32, b: f32, c: f32) -> f32 { a.max(b.max(c)) @@ -92,16 +119,22 @@ fn score_max_ternary(values: &[f32; 64], previous: f32) -> f32 { ) } +/// Flatten row, split and latent column to a usize FP32 workspace element offset. #[inline(always)] pub(crate) fn partial_offset(row: u32, split: u32, splits: u32, column: u32) -> usize { (row as usize * splits as usize + split as usize) * 512 + column as usize } +/// Pack a warp’s 32-row group and a column displacement into a live TMEM base address. #[inline(always)] fn row_address(tmem: u32, warp: u32, columns: u32) -> u32 { tmem + ((warp * 32) << 16) + columns } +/// Consume nine key chunks and publish one QK score tile, advancing persistent tickets. +/// +/// # Safety +/// The leader CTA’s issuing warp must participate with initialized rings, live shared/TMEM storage and a completed query load. #[inline(always)] unsafe fn issue_qk( smem: *mut u8, @@ -133,6 +166,10 @@ unsafe fn issue_qk( *score_ticket += 1; } +/// Consume eight value chunks and one probability buffer, then publish the updated output accumulator. +/// +/// # Safety +/// The leader CTA’s issuing warp must participate with live storage and matching ring tickets. Set first only before the first accumulation. #[inline(always)] #[allow(clippy::too_many_arguments)] // Independent pipeline tickets make ownership explicit. unsafe fn issue_pv( @@ -180,6 +217,10 @@ unsafe fn issue_pv( *o_ticket += 1; } +/// Load one page’s nine key chunks and, on the first tile, the corresponding latent and rotary queries. +/// +/// # Safety +/// Both CTAs’ load warps must supply their rank, matching live descriptors, shared storage and initialized rings. #[inline(always)] #[allow(clippy::too_many_arguments)] // Typed maps and separate pipeline lifetimes are explicit. unsafe fn load_qk( @@ -249,6 +290,10 @@ unsafe fn load_qk( } } +/// Release the consumed page-index slot and load its eight value chunks into the KV ring. +/// +/// # Safety +/// Both load warps must have waited for this page ticket and use matching rank, descriptor, storage and KV cursors. #[inline(always)] unsafe fn load_v( smem: *mut u8, @@ -293,10 +338,15 @@ unsafe fn load_v( /// Softmax's tickets span persistent work; only the row statistics reset per work item. struct SoftmaxState { + /// Next score buffer consumed; continues across persistent work items. score_ticket: u32, + /// Next probability buffer produced; also selects its maximum-exchange buffer. probability_ticket: u32, + /// Next metadata buffer published to output-correction threads. correction_ticket: u32, + /// Running unscaled QK maximum for this row and split. maximum: f32, + /// Running softmax denominator relative to the current row maximum. sum: f32, } @@ -425,6 +475,7 @@ unsafe fn softmax_tile( state.correction_ticket += 1; } +/// Persistent two-CTA decode and the optional split-reduction kernel with host launch wrappers. #[cuda_module] pub mod kernels { use super::*; diff --git a/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/src/main.rs b/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/src/main.rs index 1cb936de61..21acdeac92 100644 --- a/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/src/main.rs +++ b/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/src/main.rs @@ -3,8 +3,11 @@ * SPDX-License-Identifier: Apache-2.0 */ +//! SM100 paged MLA decode runner: validate serialized inputs, launch persistent clusters and report timing. + #![feature(f16)] +/// CLI comparison dispatch and serialized-input adapter selection. mod comparison; mod kernel; @@ -17,23 +20,36 @@ use std::{ path::{Path, PathBuf}, }; +/// Host-runner result carrying filesystem, validation or CUDA errors. type Result = std::result::Result>; +/// Serialized problem dimensions and scale factors for the fixed MLA decode profile. #[derive(Deserialize)] #[serde(deny_unknown_fields)] struct Config { + /// Number of independent decode requests. batch: u32, + /// Valid cache tokens per request; the final physical page may include padding. kv_length: u32, + /// Requested contiguous page partitions per request, including possible empty partitions. splits: u32, + /// Query heads per request; this kernel requires 128. heads: u32, + /// Query tokens per request; decode requires exactly one. query_length: u32, + /// Latent features per query and cache row; this kernel requires 512. latent: u32, + /// Rotary-position features per query and key row; this kernel requires 64. rope: u32, + /// Cache rows per physical page; this kernel requires 128. page_size: u32, + /// Positive finite multiplier applied to QK scores before softmax. softmax_scale: f32, + /// Finite multiplier applied to the normalized attention output. output_scale: f32, } impl Config { + /// Reject unsupported dimensions, invalid scales and overflowing launch or allocation sizes. fn validate(&self) -> Result<()> { if ( self.heads, @@ -82,25 +98,36 @@ impl Config { } Ok(()) } + /// Number of 128-token pages needed for one request, rounding up the tail. fn pages(&self) -> u32 { self.kv_length.div_ceil(128) } + /// Number of nonempty partitions under the kernel’s fixed pages-per-split schedule. fn active_splits(&self) -> u32 { let pages = self.pages(); pages.div_ceil(pages.div_ceil(self.splits)) } } +/// Serialized-input adapter paths, launch override and timing controls. struct Args { + /// Directory containing config.json, FP16 arrays and the physical-page table. input: PathBuf, + /// Directory receiving output.f16 and base-2 lse.f32 arrays. output: PathBuf, + /// Optional resident two-CTA cluster count; checked against capacity and available work. clusters: Option, + /// Untimed complete invocations before each measured round. warmup: usize, + /// Complete invocations averaged into each timing sample. iterations: usize, + /// Independent timing windows reported in microseconds per invocation. rounds: usize, + /// Skip timing after the initial finite-output check. validation_only: bool, } impl Args { + /// Parse adapter options and aliases, requiring input/output paths and positive launch counts. fn parse() -> Result { let mut input = None; let mut output = None; @@ -149,6 +176,7 @@ impl Args { } } +/// Read exactly count little-endian FP16 bit patterns and reject nonfinite values. fn load_u16(path: &Path, count: usize) -> Result> { let bytes = std::fs::read(path)?; if bytes.len() != count.checked_mul(2).ok_or("input byte count overflow")? { @@ -166,6 +194,7 @@ fn load_u16(path: &Path, count: usize) -> Result> { Ok(values) } +/// Read exactly count little-endian page indices, bounded by the physical cache allocation. fn load_pages(path: &Path, count: usize, physical_pages: usize) -> Result> { let bytes = std::fs::read(path)?; if bytes.len() != count.checked_mul(4).ok_or("page table overflow")? { @@ -186,6 +215,7 @@ fn load_pages(path: &Path, count: usize, physical_pages: usize) -> Result Result<()> { if comparison::dispatch(&std::env::args_os().skip(1).collect::>())? { return Ok(()); @@ -390,9 +420,11 @@ fn main() -> Result<()> { Ok(()) } +/// Host-only checks for swizzle alignment, workspace indexing and supported problem bounds. #[cfg(test)] mod tests { use super::*; + /// Every eight-element probability vector must remain contiguous and eight-element aligned. #[test] fn probability_vectors_remain_contiguous_and_aligned_under_swizzle() { for row in 0..64 { @@ -409,6 +441,7 @@ mod tests { } } + /// Workspace indexing must retain offsets beyond u32 element addressing. #[test] fn split_workspace_offset_does_not_wrap_at_four_gib_elements() { assert_eq!( @@ -416,6 +449,7 @@ mod tests { 17_179_869_183usize ); } + /// Reject physical cache row coordinates that cannot fit the TMA i32 coordinate type. #[test] fn reject_large_physical_cache_coordinates_before_allocation() { let cfg = Config { @@ -432,6 +466,7 @@ mod tests { }; assert!(cfg.validate().is_err()); } + /// Count only nonempty splits while preserving the kernel’s ceiling-based page partition. #[test] fn active_splits_ignore_empty_work_without_changing_partitioning() { let mut cfg = Config { @@ -454,6 +489,7 @@ mod tests { assert_eq!(cfg.active_splits(), expected); } } + /// Reject a zero score multiplier before allocating or launching. #[test] fn reject_nonpositive_softmax_scale() { let cfg = Config { diff --git a/cuda-oxide/scripts/check-primitives-api-compile.py b/cuda-oxide/scripts/check-primitives-api-compile.py index 87d84763e5..ed067aec62 100644 --- a/cuda-oxide/scripts/check-primitives-api-compile.py +++ b/cuda-oxide/scripts/check-primitives-api-compile.py @@ -8,6 +8,7 @@ import os from pathlib import Path import re +import struct import subprocess import tempfile import tomllib @@ -17,7 +18,7 @@ ROOT = Path(__file__).resolve().parents[1] PROFILE = "cutlass-primitives-mlir22" ARCH = "sm_100a" -KERNELS = ["primitives_grid_constant", "primitives_lifecycle", "primitives_tma", "primitives_sync", "primitives_mma"] +KERNELS = ["primitives_grid_constant", "primitives_lifecycle", "primitives_tma", "primitives_sync", "primitives_mma", "primitives_attention"] def require(condition, message): @@ -55,11 +56,116 @@ def visit(block): return visit(next(iter(blocks))) +def verify_attention_ir(lines): + """Check the CTM facade's instruction contracts at the consumer boundary.""" + text = "\n".join(lines) + definitions = {line.strip().split(" = ", 1)[0]: line for line in lines if " = " in line} + constants = { + name: (struct.unpack("!f", int(value, 16).to_bytes(4, "big"))[0] + if kind == "f32" and value.startswith("0x") else float(value)) + for name, value, kind in re.findall( + r'(%v\d+) = "arith.constant"\(\) <\{value = ([^ :]+) : (i32|i64|f32)', text) + } + + def operands(line): + args = re.search(r'"[^\"]+"\(([^)]*)\)', line).group(1) + return args.split(", ") if args else [] + + def constant_value(arg): + if arg in constants: + return constants[arg] + # Carrier constructors retain arrays in the exported, unoptimized IR. + # Follow each extract back to the constant stored at that array index. + extract = definitions[arg] + require('"llvm.extractvalue"' in extract, "expected a constant or carrier register") + index = int(re.search(r"position = array", extract).group(1)) + load = definitions[operands(extract)[0]] + require('"llvm.load"' in load, "expected carrier array load") + pointer = operands(load)[0] + initializers = [] + for line in lines: + if '"llvm.store"' not in line: + continue + value, destination = operands(line) + gep = definitions.get(destination, "") + if '"llvm.getelementptr"' in gep: + base, offset = operands(gep) + if base == pointer and constants[offset] == index: + initializers.append(value) + require(len(initializers) == 1, "carrier register must have one constant initializer") + return constants[initializers[0]] + + def values(line): + return [constant_value(arg) for arg in operands(line)] + + def one(fragment): + matches = [line for line in lines if fragment in line] + require(len(matches) == 1, f"expected exactly one attention instruction {fragment}") + return matches[0] + + one('"nvvm.cp.async.commit.group"') + waits = [line for line in lines if '"nvvm.cp.async.wait.group"' in line] + require([int(re.search(r"n = (\d+) : i32", line).group(1)) for line in waits] == [0, 7], + "cp.async wait boundary immediates changed") + for name, value in [("llvm.nvvm.ex2.approx.f32", 1.25), + ("llvm.nvvm.ex2.approx.ftz.f32", 2.5), + ("llvm.nvvm.rcp.approx.ftz.f", 4.0)]: + line = one(f'intrin = "{name}"') + require(values(line) == [value] and '(f32) -> f32' in line, + f"attention math input or exact modifier changed: {name}") + shuffle = one('intrin = "llvm.nvvm.shfl.sync.bfly.f32"') + args = operands(shuffle) + require(len(args) == 4 and int(constants[args[0]]) & 0xffffffff == 0xffffffff + and values(shuffle)[1:] == [8.0, 2.0, 31.0], + "shuffle must retain full-warp mask, XOR lane selector and clamp 31") + pack = one('intrin = "llvm.nvvm.ff2f16x2.rn"') + require(values(pack) == [2.5, 1.25] and '(f32, f32) -> vector<2xf16>' in pack, + "FP16 packing must preserve round-nearest and LLVM high/low operand order") + packed_result = pack.strip().split(" = ", 1)[0] + require(any('"llvm.bitcast"' in line and operands(line) == [packed_result] + and '(vector<2xf16>) -> i32' in line for line in lines), + "packed half pair lost its raw u32 carrier") + mma = one('"nvvm.mma.sync"') + for attr in ["#nvvm.shape", "layoutA = #nvvm.mma_layout", + "layoutB = #nvvm.mma_layout", "multiplicandAPtxType = #nvvm.mma_type", + "multiplicandBPtxType = #nvvm.mma_type", "operandSegmentSizes = array"]: + require(attr in mma, f"FP16 MMA selector changed: {attr}") + args = operands(mma) + require(len(args) == 10, "FP16 MMA operand count changed") + registers = [] + for arg in args[:6]: + cast = definitions[arg] + require('"llvm.bitcast"' in cast and '(i32) -> vector<2xf16>' in cast, + "FP16 MMA packed register must be bitcast, never numerically converted") + registers.append(int(constant_value(operands(cast)[0]))) + require(registers == [0x3c003c00, 0x40004000, 0x42004200, 0x44004400, 0x46004600, 0x48004800] + and [constant_value(arg) for arg in args[6:]] == [1.0, 2.0, 3.0, 4.0], + "FP16 MMA must preserve A[4], B[2], C[4] register order") + aggregate = mma.strip().split(" = ", 1)[0] + extracts = [line for line in lines if '"llvm.extractvalue"' in line and operands(line) == [aggregate]] + require([int(re.search(r"position = array", line).group(1)) for line in extracts] == [0, 1, 2, 3], + "FP16 MMA result register order changed") + for index, extract in enumerate(extracts): + result = extract.strip().split(" = ", 1)[0] + inserts = [line for line in lines if '"llvm.insertvalue"' in line + and operands(line)[-1] == result] + require(len(inserts) == 1 and f"position = array" in inserts[0], + "FP16 MMA result adapter reordered the accumulator carrier") + instructions = re.findall(r'"((?:nvvm|llvm\.nvvm)\.[^"]+)"', text) + require(instructions == ["nvvm.cp.async.commit.group", "nvvm.cp.async.wait.group", + "nvvm.cp.async.wait.group", "llvm.nvvm.ex2.approx.f32", + "llvm.nvvm.ex2.approx.ftz.f32", "llvm.nvvm.rcp.approx.ftz.f", + "llvm.nvvm.shfl.sync.bfly.f32", "llvm.nvvm.ff2f16x2.rn", "nvvm.mma.sync"], + "attention probe acquired hidden instructions or changed issue order") + return instructions + + def verify_ir(text, log, evidence): require(not re.search(r'"cute\.|!cute\.', text), "primitives input contains CuTe operations/types") functions = probe_functions(text) require(set(KERNELS) <= functions.keys(), "missing primitive probe function") paths = {name: instruction_path(functions[name]) for name in KERNELS} + attention = verify_attention_ir(paths["primitives_attention"]) events = lambda lines: re.findall(r'"(nvvm\.[^"]+)"', "\n".join(lines)) lifecycle = events(paths["primitives_lifecycle"]) require(lifecycle == ["nvvm.tcgen05.alloc", "nvvm.tcgen05.relinquish_alloc_permit", "nvvm.tcgen05.dealloc", "nvvm.mbarrier.init", "nvvm.mbarrier.inval"], f"unexpected lifecycle instructions: {lifecycle}") @@ -125,7 +231,9 @@ def verify_ir(text, log, evidence): counts = Counter() for (_, marker), count in pairs(after).items(): counts[marker] += count - for marker, count in {"v1:i0359": 1, "v1:i0360": 1, "v1:i0763": 8, "v1:i0664": 1, "v1:i0357": 1, "v1:i0331": 3, "v1:i0336": 1, "v1:i0308": 1, "v1:i0100": 1, "v1:i0342": 2, "v1:i0367": 1, "v1:i0887": 1}.items(): + for marker, count in {"v1:i0359": 1, "v1:i0360": 1, "v1:i0763": 8, "v1:i0664": 1, "v1:i0357": 1, "v1:i0331": 3, "v1:i0336": 1, "v1:i0308": 1, "v1:i0100": 1, "v1:i0342": 2, "v1:i0367": 1, "v1:i0887": 1, + "v1:i0094": 1, "v1:i0096": 2, "v1:i0819": 1, "v1:i0820": 1, + "v1:i0788": 1, "v1:i0055": 1, "v1:i0081": 1, "v1:i0106": 1}.items(): require(counts[marker] == count, f"incorrect actual Rust importer count for {marker}: {counts[marker]}") for section in [before, after]: line = next(line for line in section.splitlines() if '"v1:i0664"' in line) @@ -133,7 +241,9 @@ def verify_ir(text, log, evidence): require(len(results.split(", ")) == 32 and line.split(" -> ", 1)[1].count("builtin.integer ui32") == 32, "raw TMEM result widths/count changed") require(f"mir.construct_array ({results})" in section, "Rust TMEM result adapter reordered registers") require("cuda_oxide_intrinsic_marker" not in text, "source markers leaked into consumer IR") - return {name: events(lines) for name, lines in paths.items()} + checks = {name: events(lines) for name, lines in paths.items()} + checks["primitives_attention"] = attention + return checks def verify_gemm_ir(text, log, evidence, name): From ab462fd927fdaa871765580b3606afc4480e388e Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Wed, 14 Oct 2026 06:51:50 +0000 Subject: [PATCH 35/99] Route SM100 CuTe attention primitives through CTM Use CTM for explicit instruction, storage, and launch APIs in SM100 prefill, decode, and their shared CuTe attention helpers. Add the missing packed FP32, TMEM transfer, barrier, copy, warp, and register-control primitives while retaining structured CuTe MMA/TMA lowering. Preserve the canonical unroll marker identity and guard the source boundary with regression tests. Update the example documentation and archive migration evidence separately from historical measurements. Both SM100a builds, all 17 B200 correctness cases, memcheck/synccheck, CTM/CuTe tests, clippy, and rustdoc pass. Decode kernel text is unchanged; prefill adds two teardown instructions with no measurable regression in the sustained dense profile checked. Signed-off-by: nihalpasham --- cuda-oxide/cuteir/README.md | 9 +- cuda-oxide/cuteir/ctm-rs/README.md | 10 + .../cuteir/ctm-rs/docs/primitives-coverage.md | 14 +- cuda-oxide/cuteir/ctm-rs/src/async_copy.rs | 24 + cuda-oxide/cuteir/ctm-rs/src/f32x2.rs | 50 + cuda-oxide/cuteir/ctm-rs/src/float.rs | 7 + cuda-oxide/cuteir/ctm-rs/src/lib.rs | 57 +- cuda-oxide/cuteir/ctm-rs/src/mbarrier.rs | 102 ++ cuda-oxide/cuteir/ctm-rs/src/memory.rs | 3 +- cuda-oxide/cuteir/ctm-rs/src/tcgen05/mod.rs | 23 + cuda-oxide/cuteir/ctm-rs/src/tcgen05/tmem.rs | 185 +++ cuda-oxide/cuteir/ctm-rs/src/warp.rs | 30 + .../cuteir/ctm-rs/tests/attention_api.rs | 13 +- .../primitives_register_count_alignment.rs | 10 + ...primitives_register_count_alignment.stderr | 44 + .../primitives_register_count_range.rs | 10 + .../primitives_register_count_range.stderr | 44 + .../primitives_register_count_runtime.rs | 12 + .../primitives_register_count_runtime.stderr | 7 + .../primitives_tmem_carrier_width.rs | 12 + .../primitives_tmem_carrier_width.stderr | 15 + .../sm100_attention_primitives.rs | 80 ++ cuda-oxide/cuteir/cute-rs/src/arch.rs | 2 +- .../cute-rs/src/sm100_attention/__compiler.rs | 2 +- .../cuteir/cute-rs/src/sm100_attention/mod.rs | 2 + .../cute-rs/src/sm100_attention/pipeline.rs | 2 +- .../cuteir/cute-rs/src/sm100_attention/tma.rs | 2 +- .../cute-rs/src/sm100_attention/tmem.rs | 45 +- .../cuteir/cute-rs/tests/sm100_attention.rs | 25 + .../sm100_fmha_prefill_cute/README.md | 21 +- .../reference/README.md | 4 +- .../reference/measurement-notes.md | 21 + .../results/ctm-migration/baseline-dense.json | 154 +++ .../correctness/both-tails-gqa.json | 87 ++ .../correctness/kv-ring-wrap.json | 87 ++ .../correctness/persistent-reuse.json | 87 ++ .../ctm-migration/correctness/query-tail.json | 87 ++ .../correctness/reference-smoke.json | 87 ++ .../correctness/short-causal.json | 87 ++ .../ctm-migration/correctness/single-key.json | 87 ++ .../correctness/unequal-causal.json | 87 ++ .../results/ctm-migration/final-dense.json | 154 +++ .../results/ctm-migration/memcheck.json | 87 ++ .../results/ctm-migration/memcheck.log | 40 + .../results/ctm-migration/opcode-counts.json | 244 ++++ .../results/ctm-migration/summary.json | 341 +++++ .../results/ctm-migration/synccheck.json | 87 ++ .../results/ctm-migration/synccheck.log | 40 + .../sm100_fmha_prefill_cute/src/kernel.rs | 73 +- .../sm100_mla_decode_fp16_cute/README.md | 21 +- .../reference/README.md | 2 + .../reference/measurement-notes.md | 21 +- .../results/ctm-migration/correctness.json | 1193 +++++++++++++++++ .../results/ctm-migration/memcheck.log | 2 + .../results/ctm-migration/opcode-counts.json | 396 ++++++ .../results/ctm-migration/sanitizers.json | 120 ++ .../results/ctm-migration/summary.json | 337 +++++ .../results/ctm-migration/synccheck.log | 2 + .../sm100_mla_decode_fp16_cute/src/kernel.rs | 69 +- 59 files changed, 4855 insertions(+), 111 deletions(-) create mode 100644 cuda-oxide/cuteir/ctm-rs/src/f32x2.rs create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_alignment.rs create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_alignment.stderr create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_range.rs create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_range.stderr create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_runtime.rs create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_runtime.stderr create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.rs create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.stderr create mode 100644 cuda-oxide/cuteir/ctm-rs/tests/compile_pass/sm100_attention_primitives.rs create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/baseline-dense.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/both-tails-gqa.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/kv-ring-wrap.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/persistent-reuse.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/query-tail.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/reference-smoke.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/short-causal.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/single-key.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/unequal-causal.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/final-dense.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/memcheck.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/memcheck.log create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/opcode-counts.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/summary.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/synccheck.json create mode 100644 cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/synccheck.log create mode 100644 cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/results/ctm-migration/correctness.json create mode 100644 cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/results/ctm-migration/memcheck.log create mode 100644 cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/results/ctm-migration/opcode-counts.json create mode 100644 cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/results/ctm-migration/sanitizers.json create mode 100644 cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/results/ctm-migration/summary.json create mode 100644 cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/reference/results/ctm-migration/synccheck.log diff --git a/cuda-oxide/cuteir/README.md b/cuda-oxide/cuteir/README.md index 1ff09d9492..e35aea024a 100644 --- a/cuda-oxide/cuteir/README.md +++ b/cuda-oxide/cuteir/README.md @@ -74,8 +74,13 @@ Each example vendors its CUTLASS source and helpers under `reference/`, checks both outputs against an independent oracle, and compares complete invocations with the same event-timing policy. Kernel names distinguish `_cute_rs` and `_cutedsl`. See the example READMEs for supported shape profiles and validation -commands. These use CuTe layouts plus NVVM operations; `cute_ext` GQA examples -belong to the separate LIR programming model. +commands. These use `cute_rs` for structured layouts, MMA and TMA, and `ctm_rs` +for explicit instructions, pipeline synchronization, TMEM transfers and packed +arithmetic. Kernel macros and storage types are also exposed through CTM; +`cuda-device` remains a direct example dependency for macro expansion. CuTe +operations keep their CuTe dialect paths under `cutlass-mlir`; the attention +examples do not use the primitives-only profile. `cute_ext` GQA examples belong +to the separate LIR programming model. The RTX 5090 attention examples target **SM120a**. They use `cute_rs` for layouts and copies, and `ctm_rs` for instruction-level operations: diff --git a/cuda-oxide/cuteir/ctm-rs/README.md b/cuda-oxide/cuteir/ctm-rs/README.md index 8cd590210b..b5ab75e692 100644 --- a/cuda-oxide/cuteir/ctm-rs/README.md +++ b/cuda-oxide/cuteir/ctm-rs/README.md @@ -43,6 +43,16 @@ instruction subset, safety responsibilities, compiler profile, and limitations. Examples are the [SM120 block-scaled GEMM](../examples/blockscale_gemm_ctm) and [SM100 FP16 GEMM](../examples/fp16_gemm_256x352_ctm). +CTM can also supply explicit instructions inside a CuTe kernel. The +[SM100 prefill](../examples/sm100_fmha_prefill_cute) and +[MLA decode](../examples/sm100_mla_decode_fp16_cute) examples retain structured +CuTe layouts, MMA and TMA while using CTM for pipeline synchronization, TMEM +transfers, packed arithmetic and execution control. These use `cutlass-mlir`, +which accepts CuTe and NVVM operations together. Their CuTe operations do not +all lower through Rust CTM calls. Kernel macros and canonical device types are +reexported by CTM; callers still need `cuda-device` in their manifest because +macro expansions refer to it internally. + ## Validation From `cuda-oxide/` with its pinned nightly: diff --git a/cuda-oxide/cuteir/ctm-rs/docs/primitives-coverage.md b/cuda-oxide/cuteir/ctm-rs/docs/primitives-coverage.md index 1fb5e634f5..588fafcac4 100644 --- a/cuda-oxide/cuteir/ctm-rs/docs/primitives-coverage.md +++ b/cuda-oxide/cuteir/ctm-rs/docs/primitives-coverage.md @@ -1,8 +1,9 @@ # Rust CUTLASS primitives `ctm_rs` is a Rust API over the existing NVVM instruction dialect, with no CuTe -layout algebra or Python dependency. It covers two-CTA SM100 FP16 GEMM and SM120 -MXFP4 GEMM, not full PTX parity. +layout algebra or Python dependency. It covers two-CTA SM100 FP16 GEMM, SM120 +MXFP4 GEMM and explicit instruction paths used by attention kernels; it does +not provide full PTX parity. - **Initial consumer:** `cutlass-primitives-mlir22`, pinned to CUTLASS compiler 4.7.0 / LLVM 22 / `sm_100a` or `sm_120a`; Pliron `06d99a3df938394583715aacf5eadb65b4791c09`. @@ -12,6 +13,15 @@ MXFP4 GEMM, not full PTX parity. ## Coverage +The primitives-only profile described below excludes CuTe operations. CTM is +also usable from the full `cutlass-mlir` profile alongside CuTe: the +[SM100 prefill](../../examples/sm100_fmha_prefill_cute/README.md) and +[MLA decode](../../examples/sm100_mla_decode_fp16_cute/README.md) examples use +CTM for explicit pipeline, TMEM, numeric and synchronization instructions, +while structured layouts, MMA and TMA retain their CuTe dialect paths. This +integration does not establish that every CuTe operation lowers through Rust +CTM or that an attention kernel is accepted by the primitives-only profile. + | Example | Covered path | | --- | --- | | [SM100 FP16](../../examples/fp16_gemm_256x352_ctm/README.md) | Two-CTA TMEM GEMM | diff --git a/cuda-oxide/cuteir/ctm-rs/src/async_copy.rs b/cuda-oxide/cuteir/ctm-rs/src/async_copy.rs index 6fad77366d..0c7fb3e2a1 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/async_copy.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/async_copy.rs @@ -36,3 +36,27 @@ 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(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(barrier: *mut crate::mbarrier::Barrier) { + unsafe { cuda_device::async_copy::cp_async_mbarrier_arrive_noinc(barrier) } +} diff --git a/cuda-oxide/cuteir/ctm-rs/src/f32x2.rs b/cuda-oxide/cuteir/ctm-rs/src/f32x2.rs new file mode 100644 index 0000000000..9630171cca --- /dev/null +++ b/cuda-oxide/cuteir/ctm-rs/src/f32x2.rs @@ -0,0 +1,50 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! SM100 packed two-lane FP32 register arithmetic. +//! +//! Values retain the canonical `u64` register carrier. Pack/unpack preserve +//! scalar bits; arithmetic requires an SM100 profile supporting native f32x2. + +/// Pack two FP32 scalar bit patterns into one register pair. +#[must_use] +#[inline(always)] +pub fn pack_f32x2(lo: f32, hi: f32) -> u64 { + cuda_device::f32x2::pack_f32x2(lo, hi) +} + +/// Unpack the low FP32 lane without changing its bits. +#[must_use] +#[inline(always)] +pub fn unpack_f32x2_lo(value: u64) -> f32 { + cuda_device::f32x2::unpack_f32x2_lo(value) +} + +/// Unpack the high FP32 lane without changing its bits. +#[must_use] +#[inline(always)] +pub fn unpack_f32x2_hi(value: u64) -> f32 { + cuda_device::f32x2::unpack_f32x2_hi(value) +} + +/// Add both FP32 lanes with round-to-nearest-even. +#[must_use] +#[inline(always)] +pub fn add_f32x2(a: u64, b: u64) -> u64 { + cuda_device::f32x2::add_f32x2(a, b) +} + +/// Multiply both FP32 lanes with round-to-nearest-even. +#[must_use] +#[inline(always)] +pub fn mul_f32x2(a: u64, b: u64) -> u64 { + cuda_device::f32x2::mul_f32x2(a, b) +} + +/// Compute fused `a * b + c` in both FP32 lanes with round-to-nearest-even. +#[must_use] +#[inline(always)] +pub fn fma_f32x2(a: u64, b: u64, c: u64) -> u64 { + cuda_device::f32x2::fma_f32x2(a, b, c) +} diff --git a/cuda-oxide/cuteir/ctm-rs/src/float.rs b/cuda-oxide/cuteir/ctm-rs/src/float.rs index 9f90c280b5..043a4102b9 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/float.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/float.rs @@ -27,3 +27,10 @@ pub fn ex2_approx_ftz_f32(value: f32) -> f32 { pub fn rcp_approx_ftz_f32(value: f32) -> f32 { cuda_device::float::rcp_approx_ftz_f32(value) } + +/// Compute `log2(value)` 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) +} diff --git a/cuda-oxide/cuteir/ctm-rs/src/lib.rs b/cuda-oxide/cuteir/ctm-rs/src/lib.rs index 2f0d8c4e18..3d1816a838 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/lib.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/lib.rs @@ -33,11 +33,12 @@ extern crate std; pub mod host; /// Kernel macros using cuda-device's launch ABI; require a direct `cuda-device` dependency. -pub use cuda_device::{cuda_module, kernel, launch_bounds, launch_contract}; +pub use cuda_device::{cluster_launch, cuda_module, kernel, launch_bounds, launch_contract}; pub mod async_copy; pub mod cluster; pub mod convert; +pub mod f32x2; pub mod float; pub mod groups; pub mod mbarrier; @@ -50,6 +51,7 @@ 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. +// Compile-time markers are reexported to preserve their defining identity. /// Hardware thread/block coordinates, without inferred thread mapping. pub mod thread { @@ -64,6 +66,59 @@ pub mod thread { } )*}; } + /// Request unrolling by placing this marker at the start of a counted `while` body. + /// `FACTOR` is zero for full unrolling, or `2..=1024` for partial unrolling. + /// This reexport preserves the defining identity used by the compiler's unroll pass. + /// Range-based `for` loops are not recognized by the current unroll pass. + pub use cuda_device::thread::__unroll_config; + + /// 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 the warpgroup's register budget to `N` (sm_90a or supported newer profile). + /// `N` is an instruction immediate: a multiple of eight in `24..=256`. + /// + /// # Safety + /// Every warpgroup thread must execute uniformly with the same count. The + /// increase must obey the register redistribution protocol and launch budget. + #[inline(always)] + pub unsafe fn setmaxnreg_inc() + where + (): ValidRegisterCount, + { + unsafe { cuda_device::thread::setmaxnreg_inc::() } + } + + /// Decrease the warpgroup's register budget to `N` (sm_90a or supported newer profile). + /// `N` is an instruction immediate: a multiple of eight in `24..=256`. + /// + /// # Safety + /// Every warpgroup thread must execute uniformly with the same count. The + /// decrease must obey the register redistribution protocol and launch budget. + #[inline(always)] + pub unsafe fn setmaxnreg_dec() + where + (): ValidRegisterCount, + { + unsafe { cuda_device::thread::setmaxnreg_dec::() } + } + coordinate!( blockDim_x, blockDim_y, diff --git a/cuda-oxide/cuteir/ctm-rs/src/mbarrier.rs b/cuda-oxide/cuteir/ctm-rs/src/mbarrier.rs index af5b35f3b4..4422e3c37d 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/mbarrier.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/mbarrier.rs @@ -139,3 +139,105 @@ pub unsafe fn try_wait_parity_cluster( ) -> 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) } +} + +/// Acquire cluster-shared async-proxy writes through the generic proxy (sm_90+). +/// +/// # 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_generic_acquire_shared_cluster_cluster() { + 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+). +/// +/// # 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_generic_release_shared_cta_cluster() { + unsafe { cuda_device::barrier::fence_proxy_async_generic_release_shared_cta_cluster() } +} + +/// Acquire cluster-shared memory at cluster scope (sm_90+), excluding global memory. +/// +/// # 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_acquire_shared_cluster_cluster() { + unsafe { cuda_device::barrier::fence_acquire_shared_cluster_cluster() } +} diff --git a/cuda-oxide/cuteir/ctm-rs/src/memory.rs b/cuda-oxide/cuteir/ctm-rs/src/memory.rs index 12399cd24b..1baf6b16b1 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/memory.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/memory.rs @@ -10,7 +10,8 @@ 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. -pub use cuda_device::{DisjointSlice, DynamicSharedArray}; +pub use cuda_device::vector::U32x4; +pub use cuda_device::{CuSimd, DisjointSlice, DynamicSharedArray}; /// Restricts physical storage formats to the instruction layouts reviewed here. mod sealed { diff --git a/cuda-oxide/cuteir/ctm-rs/src/tcgen05/mod.rs b/cuda-oxide/cuteir/ctm-rs/src/tcgen05/mod.rs index 7ee782ecb2..9cab2d865c 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/tcgen05/mod.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/tcgen05/mod.rs @@ -4,6 +4,9 @@ */ //! 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. The crate's `docs/primitives-coverage.md` gives a complete @@ -50,3 +53,23 @@ pub unsafe fn commit_multicast_cg2( cuda_device::tcgen05::tcgen05_commit_multicast_cg2(barrier.as_raw().cast(), mask.bits()) } } + +/// Release the group-1 tensor-memory allocation permit (SM100). +/// One full warp in the CTA must execute uniformly. All tcgen05 operations in +/// the kernel must use group-1 mode; allocation/release synchronization remains explicit. +#[inline(always)] +pub fn relinquish_alloc_permit() { + cuda_device::tcgen05::tcgen05_relinquish_alloc_permit() +} + +/// 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(barrier: super::memory::SharedAddress<'_, crate::mbarrier::Barrier>) { + unsafe { cuda_device::tcgen05::tcgen05_commit(barrier.as_raw().cast()) } +} diff --git a/cuda-oxide/cuteir/ctm-rs/src/tcgen05/tmem.rs b/cuda-oxide/cuteir/ctm-rs/src/tcgen05/tmem.rs index 090ebaaa1f..c043e99b04 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/tcgen05/tmem.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/tcgen05/tmem.rs @@ -3,6 +3,10 @@ * 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, TmemDeallocated, TmemReady, TmemUninit}; @@ -275,6 +279,187 @@ where } } +/// 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. +/// All 32 warp lanes must execute convergently. No CTA barrier, producer wait +/// or memory/proxy fence is inserted. +#[inline(always)] +pub fn load_wait() { + cuda_device::tcgen05::tcgen05_load_wait() +} + +/// Complete the executing thread's previously issued SM100 TMEM stores. +/// All 32 warp lanes must execute convergently. No CTA barrier, consumer +/// synchronization or memory/proxy fence is inserted. +#[inline(always)] +pub fn store_wait() { + cuda_device::tcgen05::tcgen05_store_wait() +} + /// Regression tests for this module’s layout, representation and rejection contracts. #[cfg(test)] mod tests { diff --git a/cuda-oxide/cuteir/ctm-rs/src/warp.rs b/cuda-oxide/cuteir/ctm-rs/src/warp.rs index 5a0179dc1f..1de72bcd47 100644 --- a/cuda-oxide/cuteir/ctm-rs/src/warp.rs +++ b/cuda-oxide/cuteir/ctm-rs/src/warp.rs @@ -34,3 +34,33 @@ pub fn is_elected_sync(mask: u32) -> bool { pub fn warp_id() -> u32 { cuda_device::warp::warp_id() } + +/// Return whether every active lane's predicate is true, using the active mask. +/// Participating lanes must execute the same vote instruction convergently. +#[inline(always)] +pub fn all(predicate: bool) -> bool { + cuda_device::warp::all(predicate) +} + +/// Synchronize the lanes named by `mask` and order their memory accesses. +/// Every non-exited lane named in the mask must execute with the same mask; +/// each calling lane must be included in it. +#[inline(always)] +pub fn sync_mask(mask: u32) { + cuda_device::warp::sync_mask(mask) +} + +/// Exchange a register with the lane selected by XORing `lane_mask`. +/// All 32 lanes must participate convergently with the same lane mask. +#[inline(always)] +pub fn shuffle_xor(value: u32, lane_mask: u32) -> u32 { + cuda_device::warp::shuffle_xor(value, lane_mask) +} + +/// Exchange a register with `src_lane` among lanes named by `mask`. +/// Every non-exited lane in the mask must execute with the same mask; the +/// source lane must participate. This inserts no memory synchronization. +#[inline(always)] +pub fn shuffle_sync(mask: u32, value: u32, src_lane: u32) -> u32 { + cuda_device::warp::shuffle_sync(mask, value, src_lane) +} diff --git a/cuda-oxide/cuteir/ctm-rs/tests/attention_api.rs b/cuda-oxide/cuteir/ctm-rs/tests/attention_api.rs index 08e3acbedd..3507d3b0ea 100644 --- a/cuda-oxide/cuteir/ctm-rs/tests/attention_api.rs +++ b/cuda-oxide/cuteir/ctm-rs/tests/attention_api.rs @@ -7,5 +7,16 @@ /// Compile the attention primitive probe with typed instruction operands and canonical storage. #[test] fn attention_primitives_accept_instruction_operands_and_canonical_storage() { - trybuild::TestCases::new().pass("tests/compile_pass/attention_primitives.rs"); + 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/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_alignment.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_alignment.rs new file mode 100644 index 0000000000..149351a384 --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_alignment.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_alignment.stderr new file mode 100644 index 0000000000..e599980ea1 --- /dev/null +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_alignment.stderr @@ -0,0 +1,44 @@ +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() + | where + | (): ValidRegisterCount, + | ^^^^^^^^^^^^^^^^^^^^^ required by this bound in `setmaxnreg_dec` diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_range.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_range.rs new file mode 100644 index 0000000000..b70838154c --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_range.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_range.stderr new file mode 100644 index 0000000000..b4099d8753 --- /dev/null +++ b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_range.stderr @@ -0,0 +1,44 @@ +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() + | where + | (): ValidRegisterCount, + | ^^^^^^^^^^^^^^^^^^^^^ required by this bound in `setmaxnreg_inc` diff --git a/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_runtime.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_runtime.rs new file mode 100644 index 0000000000..541dee1bb8 --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_runtime.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_runtime.stderr new file mode 100644 index 0000000000..af7e13a5a0 --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.rs new file mode 100644 index 0000000000..822af8e9bb --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.stderr b/cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.stderr new file mode 100644 index 0000000000..d9c28e9d4b --- /dev/null +++ b/cuda-oxide/cuteir/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/cuteir/ctm-rs/tests/compile_pass/sm100_attention_primitives.rs b/cuda-oxide/cuteir/ctm-rs/tests/compile_pass/sm100_attention_primitives.rs new file mode 100644 index 0000000000..3c06515c8a --- /dev/null +++ b/cuda-oxide/cuteir/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, 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); + mbarrier::fence_proxy_async_generic_acquire_shared_cluster_cluster(); + mbarrier::fence_proxy_async_generic_release_shared_cta_cluster(); + mbarrier::fence_acquire_shared_cluster_cluster(); + async_copy::cp_async_ca_4(shared, global); + async_copy::cp_async_mbarrier_arrive_noinc(bar); + tcgen05::commit(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(); + } + 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 _ = warp::all(true); + let _ = warp::shuffle_xor(0, 1); + let _ = warp::shuffle_sync(u32::MAX, 0, 0); + warp::sync_mask(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/cuteir/cute-rs/src/arch.rs b/cuda-oxide/cuteir/cute-rs/src/arch.rs index 9ea21cc2ac..a62f69dca9 100644 --- a/cuda-oxide/cuteir/cute-rs/src/arch.rs +++ b/cuda-oxide/cuteir/cute-rs/src/arch.rs @@ -19,5 +19,5 @@ /// it does not check or repair a divergent value. #[inline(always)] pub unsafe fn assume_warp_uniform(value: u32) -> u32 { - cuda_device::warp::shuffle_sync(u32::MAX, value, 0) + ctm_rs::warp::shuffle_sync(u32::MAX, value, 0) } diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/__compiler.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/__compiler.rs index 9dbc602604..1d5ff9056a 100644 --- a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/__compiler.rs +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/__compiler.rs @@ -6,7 +6,7 @@ //! Exact compiler markers for independent attention operations. use crate::TmaDesc; -use cuda_device::Barrier; +use ctm_rs::mbarrier::Barrier; /// Device-only shared/shared MMA marker; operand displacements count bytes from retained bases. /// /// # Safety diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/mod.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/mod.rs index d573b5e3b2..6c0bee39f6 100644 --- a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/mod.rs +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/mod.rs @@ -7,6 +7,8 @@ //! //! 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; diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/pipeline.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/pipeline.rs index ab9d5b5071..62580abec5 100644 --- a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/pipeline.rs +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/pipeline.rs @@ -25,8 +25,8 @@ use core::marker::PhantomData; use core::mem::{align_of, size_of}; +use ctm_rs::mbarrier::{self as barrier, Barrier}; use ctm_rs::{cluster, groups::CtaGroup2, memory::SharedAddress, tcgen05}; -use cuda_device::barrier::{self, Barrier}; /// Restricts edge policies and cursor roles to the implementations in this module. mod sealed { diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/tma.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/tma.rs index e67f6afb21..986ac941cd 100644 --- a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/tma.rs +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/tma.rs @@ -8,7 +8,7 @@ use super::{__compiler, SharedTile, SmemLayout}; use crate::TmaDesc; use core::marker::PhantomData; -use cuda_device::Barrier; +use ctm_rs::mbarrier::Barrier; /// A rank-three TMA copy into the same physical layout used by MMA. pub struct TmaCopy( diff --git a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/tmem.rs b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/tmem.rs index d5ae2be419..3715d66471 100644 --- a/cuda-oxide/cuteir/cute-rs/src/sm100_attention/tmem.rs +++ b/cuda-oxide/cuteir/cute-rs/src/sm100_attention/tmem.rs @@ -3,9 +3,10 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! Warp-collective register fragments for tensor-memory loads and stores. +//! Warp-collective register fragments for CTM tensor-memory loads and stores. +//! Conversion and chunking preserve the explicit completion choices of each API. -use cuda_device::CuSimd; +use ctm_rs::{memory::CuSimd, tcgen05::tmem, thread}; /// Warp-collective 32-data-path TMEM register fragment, one row per lane. pub struct TmemFragment { @@ -23,7 +24,7 @@ macro_rules! fragment { #[inline(always)] pub unsafe fn load(address: u32) -> Self { let result = unsafe { Self::load_deferred(address) }; - cuda_device::tcgen05::tcgen05_load_wait(); + tmem::load_wait(); result } /// Load without a completion wait, retaining register dependencies. @@ -35,11 +36,11 @@ macro_rules! fragment { /// 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 { cuda_device::tcgen05::$load(address) }.to_array(); + let raw = unsafe { tmem::$load(address) }.to_array(); let mut values = [0.0; $n]; let mut i = 0; while i < $n { - cuda_device::thread::__unroll_config::<0>(); + thread::__unroll_config::<0>(); values[i] = f32::from_bits(raw[i]); i += 1; } @@ -52,7 +53,7 @@ macro_rules! fragment { #[inline(always)] pub unsafe fn store(&self, address: u32) { unsafe { self.store_deferred(address) }; - cuda_device::tcgen05::tcgen05_store_wait(); + tmem::store_wait(); } /// Store without waiting, allowing disjoint TMEM chunks to overlap. /// # Safety @@ -64,7 +65,7 @@ macro_rules! fragment { let mut raw = [0u32; $n]; let mut i = 0; while i < $n { - cuda_device::thread::__unroll_config::<0>(); + thread::__unroll_config::<0>(); raw[i] = self.values[i].to_bits(); i += 1; } @@ -76,7 +77,7 @@ macro_rules! fragment { #[inline(always)] pub unsafe fn store_packed(address: u32, raw: [u32; $n]) { unsafe { Self::store_packed_deferred(address, raw) }; - cuda_device::tcgen05::tcgen05_store_wait(); + tmem::store_wait(); } /// Store packed words without waiting for completion. /// # Safety @@ -84,16 +85,16 @@ macro_rules! fragment { /// `store_deferred`. #[inline(always)] pub unsafe fn store_packed_deferred(address: u32, raw: [u32; $n]) { - unsafe { cuda_device::tcgen05::$store(address, CuSimd::new(raw)) }; + unsafe { tmem::$store(address, CuSimd::new(raw)) }; } } }; } -fragment!(2, tcgen05_ld_32x32b_x2_raw, tcgen05_st_32x32b_x2_raw); -fragment!(4, tcgen05_ld_32x32b_x4_raw, tcgen05_st_32x32b_x4_raw); -fragment!(8, tcgen05_ld_32x32b_x8_raw, tcgen05_st_32x32b_x8_raw); -fragment!(16, tcgen05_ld_32x32b_x16_raw, tcgen05_st_32x32b_x16_raw); -fragment!(32, tcgen05_ld_32x32b_x32_raw, tcgen05_st_32x32b_x32_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. @@ -108,7 +109,7 @@ macro_rules! chunked_fragment { #[inline(always)] pub unsafe fn load(address: u32) -> Self { let result = unsafe { Self::load_deferred(address) }; - cuda_device::tcgen05::tcgen05_load_wait(); + tmem::load_wait(); result } /// Load x32 chunks without a completion wait. @@ -119,13 +120,13 @@ macro_rules! chunked_fragment { /// required before-thread-sync fence. #[inline(always)] pub unsafe fn load_deferred(address: u32) -> Self { - $(let $chunk = unsafe { cuda_device::tcgen05::tcgen05_ld_32x32b_x32_raw(address + $offset) };)+ + $(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 { - cuda_device::thread::__unroll_config::<0>(); + thread::__unroll_config::<0>(); values[$offset + i] = f32::from_bits(raw[i]); i += 1; } @@ -139,7 +140,7 @@ macro_rules! chunked_fragment { #[inline(always)] pub unsafe fn store(&self, address: u32) { unsafe { self.store_deferred(address) }; - cuda_device::tcgen05::tcgen05_store_wait(); + tmem::store_wait(); } /// Store x32 chunks without waiting for completion. /// # Safety @@ -151,7 +152,7 @@ macro_rules! chunked_fragment { let mut raw = [0; $n]; let mut i = 0; while i < $n { - cuda_device::thread::__unroll_config::<0>(); + thread::__unroll_config::<0>(); raw[i] = self.values[i].to_bits(); i += 1; } @@ -163,7 +164,7 @@ macro_rules! chunked_fragment { #[inline(always)] pub unsafe fn store_packed(address: u32, raw: [u32; $n]) { unsafe { Self::store_packed_deferred(address, raw) }; - cuda_device::tcgen05::tcgen05_store_wait(); + tmem::store_wait(); } /// Store packed x32 chunks without waiting for completion. /// # Safety @@ -175,11 +176,11 @@ macro_rules! chunked_fragment { let mut chunk = [0; 32]; let mut i = 0; while i < 32 { - cuda_device::thread::__unroll_config::<0>(); + thread::__unroll_config::<0>(); chunk[i] = raw[$offset + i]; i += 1; } - unsafe { cuda_device::tcgen05::tcgen05_st_32x32b_x32_raw(address + $offset, CuSimd::new(chunk)) }; + unsafe { tmem::store_32x32b_x32(address + $offset, CuSimd::new(chunk)) }; )+ } } diff --git a/cuda-oxide/cuteir/cute-rs/tests/sm100_attention.rs b/cuda-oxide/cuteir/cute-rs/tests/sm100_attention.rs index d0dee730ed..e28c147bd7 100644 --- a/cuda-oxide/cuteir/cute-rs/tests/sm100_attention.rs +++ b/cuda-oxide/cuteir/cute-rs/tests/sm100_attention.rs @@ -7,6 +7,31 @@ use cute_rs::sm100_attention::{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 cuteir = 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/sm100_attention/pipeline.rs", + "cute-rs/src/sm100_attention/tmem.rs", + "cute-rs/src/sm100_attention/tma.rs", + "cute-rs/src/sm100_attention/__compiler.rs", + ]; + for path in paths { + let source = std::fs::read_to_string(cuteir.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(); diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/README.md b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/README.md index 2f2154c4c8..10211b6771 100644 --- a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/README.md +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/README.md @@ -3,6 +3,14 @@ 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 explicit instructions: pipeline synchronization, TMEM transfers, +packed arithmetic, warp/cluster coordination and register control. CTM also +exposes the kernel macros and canonical storage types. The manifest retains +`cuda-device` because the macro expansions still refer to that crate internally. +CuTe operations retain their CuTe dialect lowering through `cutlass-mlir`; +this example uses both APIs rather than the primitives-only compiler profile. + - 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 @@ -12,6 +20,12 @@ From `cuda-oxide/`, with CUTLASS 4.7 installed (`cargo oxide toolchain install c and the [Python reference environment](reference/README.md) available: ```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 @@ -34,9 +48,9 @@ Add `--json comparison.json` to save the full report. `--shape` is `B,Sq,Sk,Hq,Hkv`; add `--causal` for causal attention or `--check-only` to validate without timing. Pass `--help` after `--` for options. -Previously recorded on B200 with CUTLASS 4.7 on both sides (2026-10-12): median -CUDA-event invocation times, **10 warmups**, 100 invocations per window, five -windows, hot cache. The commands above use **100 warmups** for new runs. +Historical results from before the CTM migration, recorded on B200 with +CUTLASS 4.7 on both sides (2026-10-12): median CUDA-event invocation times, +**10 warmups**, 100 invocations per window, five windows, hot cache. The commands above use **100 warmups** for new runs. | Profile (B, Sq, Sk, Hq, Hkv) | CuTe-RS | CuTeDSL | | --- | ---: | ---: | @@ -49,6 +63,7 @@ a kernel-active speedup, and the dense persistent timing ranges overlap. With 2,000 warmups and 301 invocations per window, dense persistent medians were **289.98 µs** (CuTe-RS) and **292.13 µs** (CuTeDSL), comparable throughput. +- [CTM migration: 8 passing B200 cases, sanitizers and compiler comparison](reference/results/ctm-migration/summary.json) - [CuTeDSL reference and setup](reference/README.md) - [Validation, benchmark commands, and detailed results](reference/measurement-notes.md) - [Shared attention functionality and compiler gaps](../sm100_mla_decode_fp16_cute/reference/attention-feature-gaps.md) diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/README.md b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/README.md index b412b749bd..966de74dd6 100644 --- a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/README.md +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/README.md @@ -10,7 +10,9 @@ This is the ordinary CuTe/NVVM programming model using `@cute.jit`, TMA, `tcgen05`, TMEM, and fused online softmax for data center Blackwell. The Rust kernel is `sm100_fmha_prefill_cute_rs`; [`../compare.py`](../compare.py) invokes the pinned source through `sm100_fmha_prefill_cutedsl`. The `cute_ext` GQA -examples use a separate LIR programming model. +examples use a separate LIR programming model. The Rust kernel combines +structured CuTe operations with explicit CTM instructions; see the +[programming boundary](../README.md). ## Setup and matched comparison diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/measurement-notes.md b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/measurement-notes.md index 2b89a31dbd..5519b7b564 100644 --- a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/measurement-notes.md +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/measurement-notes.md @@ -5,6 +5,27 @@ Schedule details, numerical checks, timing methodology, and evidence for the pinned CuTeDSL source and Python environment. Commands below run from the example directory unless stated otherwise. +The historical reports linked below predate the CTM migration and describe the +source revision recorded in each artifact. Current migration evidence follows +in the next section; the reproduction commands validate and measure new builds. +Structured CuTe operations keep their existing dialect paths, while explicit +instruction calls now use CTM as described in the [example README](../README.md). + +## CTM migration validation + +The [migration report](results/ctm-migration/summary.json) records **8 passing +B200 correctness cases** against pinned CuTeDSL and the independent oracle on +2026-10-14, plus source hashes and the compiler comparison with `d642dcbfbf`. +Resources are unchanged; typed TMEM teardown adds two non-padding instructions. +Both memcheck and synccheck report zero errors. This report is separate from +the historical performance results below. +A matched dense persistent run at `(B,Sq,Sk,Hq,Hkv)=(4,4096,4096,8,8)` used +2,000 warmups, 301 invocations per window and five windows. Baseline and migrated +medians were **290.322 µs** and **290.476 µs** (+0.053%), within the observed +spread. This one-profile invocation-throughput check found no measurable +regression; it does not establish broad performance parity. Both raw timing +reports are linked from the migration summary. + ## Build and reproduce From `cuda-oxide/`: diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/baseline-dense.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/baseline-dense.json new file mode 100644 index 0000000000..c5cf349635 --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/baseline-dense.json @@ -0,0 +1,154 @@ +{ + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "source_sha256": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "source_closure_sha256": { + "cute/blackwell/kernel/attention/fmha/fmha.py": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "helpers/fmha_helpers.py": "cbe090eb5567b7cdbd8b4944cbd579084fd60379c463bcfb011e6d4b28ef20e1", + "helpers/__init__.py": "9cf1ae75e6473e5acf38b8ef7b498a14ff74fd54741a930708f6995fcacb7581" + }, + "reference_compiler_optimization_level": 2, + "sanitizer": null, + "rust_executable_sha256": "24d8c9d44e7976cc84b4df2d12ab0efcfcaef5d2a8fcc4bb5d3b609d3faa47bc", + "config": { + "batch": 4, + "query": 4096, + "keys": 4096, + "heads": 8, + "kv_heads": 8, + "head_dim": 128, + "causal": false + }, + "seed": 1111, + "pattern": "gaussian", + "oracle": { + "dtype": "float32", + "matmul_precision": "highest", + "allow_tf32": false, + "query_chunk": 128 + }, + "rust_kernel": { + "blocks": 148, + "cache": "hot", + "causal": false, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 4, + 4096, + 4096, + 8, + 8 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 1.9516795873641968e-05, + "rms": 2.3820080032120917e-06, + "rtol": 1e-05 + }, + "warmup": 10, + "windows": 3 + }, + "validation": { + "sm100_fmha_prefill_cute_rs": { + "max_abs": 1.9516795873641968e-05, + "rms": 2.3820080059522297e-06, + "atol": 0.1, + "rtol": 1e-05, + "elements": 16777216 + }, + "sm100_fmha_prefill_cutedsl": { + "max_abs": 1.9516795873641968e-05, + "rms": 2.3820159640308702e-06, + "atol": 0.1, + "rtol": 1e-05, + "elements": 16777216 + }, + "cross_implementation": { + "max_abs": 1.52587890625e-05, + "rms": 2.3589489472897185e-08, + "atol": 0.1, + "rtol": 1e-05, + "elements": 16777216 + } + }, + "timing": { + "mode": "direct-cuda-events", + "cache": "hot", + "measurement_scope": "complete invocation throughput including host submission gaps", + "warmup": 2000, + "iterations": 301, + "windows": 5, + "sm100_fmha_prefill_cutedsl": { + "samples_us": [ + 292.812537513302, + 286.8274929515547, + 289.1170210220489, + 291.8325531918345, + 290.1223014755503 + ], + "median_us": 290.1223014755503 + }, + "sm100_fmha_prefill_cute_rs": { + "blocks": 148, + "cache": "hot", + "causal": false, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 301, + "median_us": 290.3222879302066, + "registers_per_thread": 128, + "samples_us": [ + 287.5866557276526, + 290.3222879302066, + 286.4804695610984, + 290.72340144667515, + 290.6687536904978 + ], + "shape": [ + 4, + 4096, + 4096, + 8, + 8 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 1.9516795873641968e-05, + "rms": 2.3820080032120917e-06, + "rtol": 1e-05 + }, + "warmup": 2000, + "windows": 5 + }, + "gpu_before": { + "query": "--query-gpu=name,uuid,driver_version,clocks.current.sm,clocks.current.memory,power.draw,power.limit,utilization.gpu", + "output": "NVIDIA B200, GPU-a4674550-4882-a08d-60b1-5efde3790e39, 615.71.09, 1830, 4000, 193.68, 1000.00, 0", + "error": "" + }, + "gpu_after": { + "query": "--query-gpu=name,uuid,driver_version,clocks.current.sm,clocks.current.memory,power.draw,power.limit,utilization.gpu", + "output": "NVIDIA B200, GPU-a4674550-4882-a08d-60b1-5efde3790e39, 615.71.09, 1830, 4000, 551.95, 1000.00, 1", + "error": "" + } + } +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/both-tails-gqa.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/both-tails-gqa.json new file mode 100644 index 0000000000..b4a611b0ad --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/both-tails-gqa.json @@ -0,0 +1,87 @@ +{ + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "source_sha256": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "source_closure_sha256": { + "cute/blackwell/kernel/attention/fmha/fmha.py": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "helpers/fmha_helpers.py": "cbe090eb5567b7cdbd8b4944cbd579084fd60379c463bcfb011e6d4b28ef20e1", + "helpers/__init__.py": "9cf1ae75e6473e5acf38b8ef7b498a14ff74fd54741a930708f6995fcacb7581" + }, + "reference_compiler_optimization_level": 2, + "sanitizer": null, + "rust_executable_sha256": "da252cfe6151cf28e583790078cf8df7125e40e3ae0c88eebf363db76cf19e61", + "config": { + "batch": 1, + "query": 257, + "keys": 385, + "heads": 4, + "kv_heads": 2, + "head_dim": 128, + "causal": true + }, + "seed": 1111, + "pattern": "gaussian", + "oracle": { + "dtype": "float32", + "matmul_precision": "highest", + "allow_tf32": false, + "query_chunk": 128 + }, + "rust_kernel": { + "blocks": 8, + "cache": "hot", + "causal": true, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 1, + 257, + 385, + 4, + 2 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 0.0003203749656677246, + "rms": 1.9929519438751464e-05, + "rtol": 1e-05 + }, + "warmup": 10, + "windows": 3 + }, + "validation": { + "sm100_fmha_prefill_cute_rs": { + "max_abs": 0.0003203749656677246, + "rms": 1.9929519112338312e-05, + "atol": 0.1, + "rtol": 1e-05, + "elements": 131584 + }, + "sm100_fmha_prefill_cutedsl": { + "max_abs": 0.0003203749656677246, + "rms": 1.9929484551539645e-05, + "atol": 0.1, + "rtol": 1e-05, + "elements": 131584 + }, + "cross_implementation": { + "max_abs": 1.52587890625e-05, + "rms": 6.651023909398646e-08, + "atol": 0.1, + "rtol": 1e-05, + "elements": 131584 + } + }, + "timing": null +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/kv-ring-wrap.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/kv-ring-wrap.json new file mode 100644 index 0000000000..9d292005f0 --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/kv-ring-wrap.json @@ -0,0 +1,87 @@ +{ + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "source_sha256": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "source_closure_sha256": { + "cute/blackwell/kernel/attention/fmha/fmha.py": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "helpers/fmha_helpers.py": "cbe090eb5567b7cdbd8b4944cbd579084fd60379c463bcfb011e6d4b28ef20e1", + "helpers/__init__.py": "9cf1ae75e6473e5acf38b8ef7b498a14ff74fd54741a930708f6995fcacb7581" + }, + "reference_compiler_optimization_level": 2, + "sanitizer": null, + "rust_executable_sha256": "da252cfe6151cf28e583790078cf8df7125e40e3ae0c88eebf363db76cf19e61", + "config": { + "batch": 1, + "query": 512, + "keys": 1153, + "heads": 8, + "kv_heads": 2, + "head_dim": 128, + "causal": false + }, + "seed": 1111, + "pattern": "gaussian", + "oracle": { + "dtype": "float32", + "matmul_precision": "highest", + "allow_tf32": false, + "query_chunk": 128 + }, + "rust_kernel": { + "blocks": 16, + "cache": "hot", + "causal": false, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 1, + 512, + 1153, + 8, + 2 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 2.6892870664596558e-05, + "rms": 4.552741129992708e-06, + "rtol": 1e-05 + }, + "warmup": 10, + "windows": 3 + }, + "validation": { + "sm100_fmha_prefill_cute_rs": { + "max_abs": 2.6892870664596558e-05, + "rms": 4.55274130217731e-06, + "atol": 0.1, + "rtol": 1e-05, + "elements": 524288 + }, + "sm100_fmha_prefill_cutedsl": { + "max_abs": 2.6892870664596558e-05, + "rms": 4.552762220555451e-06, + "atol": 0.1, + "rtol": 1e-05, + "elements": 524288 + }, + "cross_implementation": { + "max_abs": 1.52587890625e-05, + "rms": 2.7120522716472806e-08, + "atol": 0.1, + "rtol": 1e-05, + "elements": 524288 + } + }, + "timing": null +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/persistent-reuse.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/persistent-reuse.json new file mode 100644 index 0000000000..74d3da08de --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/persistent-reuse.json @@ -0,0 +1,87 @@ +{ + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "source_sha256": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "source_closure_sha256": { + "cute/blackwell/kernel/attention/fmha/fmha.py": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "helpers/fmha_helpers.py": "cbe090eb5567b7cdbd8b4944cbd579084fd60379c463bcfb011e6d4b28ef20e1", + "helpers/__init__.py": "9cf1ae75e6473e5acf38b8ef7b498a14ff74fd54741a930708f6995fcacb7581" + }, + "reference_compiler_optimization_level": 2, + "sanitizer": null, + "rust_executable_sha256": "da252cfe6151cf28e583790078cf8df7125e40e3ae0c88eebf363db76cf19e61", + "config": { + "batch": 4, + "query": 4096, + "keys": 4096, + "heads": 8, + "kv_heads": 8, + "head_dim": 128, + "causal": false + }, + "seed": 1111, + "pattern": "gaussian", + "oracle": { + "dtype": "float32", + "matmul_precision": "highest", + "allow_tf32": false, + "query_chunk": 128 + }, + "rust_kernel": { + "blocks": 148, + "cache": "hot", + "causal": false, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 4, + 4096, + 4096, + 8, + 8 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 1.9516795873641968e-05, + "rms": 2.3820080032120917e-06, + "rtol": 1e-05 + }, + "warmup": 10, + "windows": 3 + }, + "validation": { + "sm100_fmha_prefill_cute_rs": { + "max_abs": 1.9516795873641968e-05, + "rms": 2.3820080059522297e-06, + "atol": 0.1, + "rtol": 1e-05, + "elements": 16777216 + }, + "sm100_fmha_prefill_cutedsl": { + "max_abs": 1.9516795873641968e-05, + "rms": 2.3820159640308702e-06, + "atol": 0.1, + "rtol": 1e-05, + "elements": 16777216 + }, + "cross_implementation": { + "max_abs": 1.52587890625e-05, + "rms": 2.3589489472897185e-08, + "atol": 0.1, + "rtol": 1e-05, + "elements": 16777216 + } + }, + "timing": null +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/query-tail.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/query-tail.json new file mode 100644 index 0000000000..713cd47485 --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/query-tail.json @@ -0,0 +1,87 @@ +{ + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "source_sha256": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "source_closure_sha256": { + "cute/blackwell/kernel/attention/fmha/fmha.py": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "helpers/fmha_helpers.py": "cbe090eb5567b7cdbd8b4944cbd579084fd60379c463bcfb011e6d4b28ef20e1", + "helpers/__init__.py": "9cf1ae75e6473e5acf38b8ef7b498a14ff74fd54741a930708f6995fcacb7581" + }, + "reference_compiler_optimization_level": 2, + "sanitizer": null, + "rust_executable_sha256": "da252cfe6151cf28e583790078cf8df7125e40e3ae0c88eebf363db76cf19e61", + "config": { + "batch": 1, + "query": 129, + "keys": 128, + "heads": 2, + "kv_heads": 2, + "head_dim": 128, + "causal": false + }, + "seed": 1111, + "pattern": "gaussian", + "oracle": { + "dtype": "float32", + "matmul_precision": "highest", + "allow_tf32": false, + "query_chunk": 128 + }, + "rust_kernel": { + "blocks": 2, + "cache": "hot", + "causal": false, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 1, + 129, + 128, + 2, + 2 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 7.070600986480713e-05, + "rms": 1.3476355289444303e-05, + "rtol": 1e-05 + }, + "warmup": 10, + "windows": 3 + }, + "validation": { + "sm100_fmha_prefill_cute_rs": { + "max_abs": 7.070600986480713e-05, + "rms": 1.3476354979502503e-05, + "atol": 0.1, + "rtol": 1e-05, + "elements": 33024 + }, + "sm100_fmha_prefill_cutedsl": { + "max_abs": 7.070600986480713e-05, + "rms": 1.3476354979502503e-05, + "atol": 0.1, + "rtol": 1e-05, + "elements": 33024 + }, + "cross_implementation": { + "max_abs": 0.0, + "rms": 0.0, + "atol": 0.1, + "rtol": 1e-05, + "elements": 33024 + } + }, + "timing": null +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/reference-smoke.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/reference-smoke.json new file mode 100644 index 0000000000..2141d6a677 --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/reference-smoke.json @@ -0,0 +1,87 @@ +{ + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "source_sha256": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "source_closure_sha256": { + "cute/blackwell/kernel/attention/fmha/fmha.py": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "helpers/fmha_helpers.py": "cbe090eb5567b7cdbd8b4944cbd579084fd60379c463bcfb011e6d4b28ef20e1", + "helpers/__init__.py": "9cf1ae75e6473e5acf38b8ef7b498a14ff74fd54741a930708f6995fcacb7581" + }, + "reference_compiler_optimization_level": 2, + "sanitizer": null, + "rust_executable_sha256": "da252cfe6151cf28e583790078cf8df7125e40e3ae0c88eebf363db76cf19e61", + "config": { + "batch": 1, + "query": 1024, + "keys": 1024, + "heads": 8, + "kv_heads": 8, + "head_dim": 128, + "causal": false + }, + "seed": 1111, + "pattern": "gaussian", + "oracle": { + "dtype": "float32", + "matmul_precision": "highest", + "allow_tf32": false, + "query_chunk": 128 + }, + "rust_kernel": { + "blocks": 32, + "cache": "hot", + "causal": false, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 1, + 1024, + 1024, + 8, + 8 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 2.750009298324585e-05, + "rms": 4.668578572098128e-06, + "rtol": 1e-05 + }, + "warmup": 10, + "windows": 3 + }, + "validation": { + "sm100_fmha_prefill_cute_rs": { + "max_abs": 2.750009298324585e-05, + "rms": 4.668578185373917e-06, + "atol": 0.1, + "rtol": 1e-05, + "elements": 1048576 + }, + "sm100_fmha_prefill_cutedsl": { + "max_abs": 2.750009298324585e-05, + "rms": 4.668540441343794e-06, + "atol": 0.1, + "rtol": 1e-05, + "elements": 1048576 + }, + "cross_implementation": { + "max_abs": 3.0517578125e-05, + "rms": 5.1290431457573504e-08, + "atol": 0.1, + "rtol": 1e-05, + "elements": 1048576 + } + }, + "timing": null +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/short-causal.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/short-causal.json new file mode 100644 index 0000000000..3a3fc4b70c --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/short-causal.json @@ -0,0 +1,87 @@ +{ + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "source_sha256": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "source_closure_sha256": { + "cute/blackwell/kernel/attention/fmha/fmha.py": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "helpers/fmha_helpers.py": "cbe090eb5567b7cdbd8b4944cbd579084fd60379c463bcfb011e6d4b28ef20e1", + "helpers/__init__.py": "9cf1ae75e6473e5acf38b8ef7b498a14ff74fd54741a930708f6995fcacb7581" + }, + "reference_compiler_optimization_level": 2, + "sanitizer": null, + "rust_executable_sha256": "da252cfe6151cf28e583790078cf8df7125e40e3ae0c88eebf363db76cf19e61", + "config": { + "batch": 1, + "query": 17, + "keys": 31, + "heads": 4, + "kv_heads": 2, + "head_dim": 128, + "causal": true + }, + "seed": 1111, + "pattern": "uniform", + "oracle": { + "dtype": "float32", + "matmul_precision": "highest", + "allow_tf32": false, + "query_chunk": 128 + }, + "rust_kernel": { + "blocks": 4, + "cache": "hot", + "causal": true, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 1, + 17, + 31, + 4, + 2 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.004999999888241291, + "max_abs": 0.0004506111145019531, + "rms": 8.420169064352974e-05, + "rtol": 1e-05 + }, + "warmup": 10, + "windows": 3 + }, + "validation": { + "sm100_fmha_prefill_cute_rs": { + "max_abs": 0.0004506111145019531, + "rms": 8.420169615419582e-05, + "atol": 0.005, + "rtol": 1e-05, + "elements": 8704 + }, + "sm100_fmha_prefill_cutedsl": { + "max_abs": 0.0004506111145019531, + "rms": 8.420169615419582e-05, + "atol": 0.005, + "rtol": 1e-05, + "elements": 8704 + }, + "cross_implementation": { + "max_abs": 0.0, + "rms": 0.0, + "atol": 0.005, + "rtol": 1e-05, + "elements": 8704 + } + }, + "timing": null +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/single-key.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/single-key.json new file mode 100644 index 0000000000..03b96b6e30 --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/single-key.json @@ -0,0 +1,87 @@ +{ + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "source_sha256": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "source_closure_sha256": { + "cute/blackwell/kernel/attention/fmha/fmha.py": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "helpers/fmha_helpers.py": "cbe090eb5567b7cdbd8b4944cbd579084fd60379c463bcfb011e6d4b28ef20e1", + "helpers/__init__.py": "9cf1ae75e6473e5acf38b8ef7b498a14ff74fd54741a930708f6995fcacb7581" + }, + "reference_compiler_optimization_level": 2, + "sanitizer": null, + "rust_executable_sha256": "da252cfe6151cf28e583790078cf8df7125e40e3ae0c88eebf363db76cf19e61", + "config": { + "batch": 1, + "query": 1, + "keys": 1, + "heads": 1, + "kv_heads": 1, + "head_dim": 128, + "causal": false + }, + "seed": 1111, + "pattern": "uniform", + "oracle": { + "dtype": "float32", + "matmul_precision": "highest", + "allow_tf32": false, + "query_chunk": 128 + }, + "rust_kernel": { + "blocks": 1, + "cache": "hot", + "causal": false, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 1, + 1, + 1, + 1, + 1 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.004999999888241291, + "max_abs": 0.0, + "rms": 0.0, + "rtol": 1e-05 + }, + "warmup": 10, + "windows": 3 + }, + "validation": { + "sm100_fmha_prefill_cute_rs": { + "max_abs": 0.0, + "rms": 0.0, + "atol": 0.005, + "rtol": 1e-05, + "elements": 128 + }, + "sm100_fmha_prefill_cutedsl": { + "max_abs": 0.0, + "rms": 0.0, + "atol": 0.005, + "rtol": 1e-05, + "elements": 128 + }, + "cross_implementation": { + "max_abs": 0.0, + "rms": 0.0, + "atol": 0.005, + "rtol": 1e-05, + "elements": 128 + } + }, + "timing": null +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/unequal-causal.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/unequal-causal.json new file mode 100644 index 0000000000..9b421b9143 --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/correctness/unequal-causal.json @@ -0,0 +1,87 @@ +{ + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "source_sha256": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "source_closure_sha256": { + "cute/blackwell/kernel/attention/fmha/fmha.py": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "helpers/fmha_helpers.py": "cbe090eb5567b7cdbd8b4944cbd579084fd60379c463bcfb011e6d4b28ef20e1", + "helpers/__init__.py": "9cf1ae75e6473e5acf38b8ef7b498a14ff74fd54741a930708f6995fcacb7581" + }, + "reference_compiler_optimization_level": 2, + "sanitizer": null, + "rust_executable_sha256": "da252cfe6151cf28e583790078cf8df7125e40e3ae0c88eebf363db76cf19e61", + "config": { + "batch": 1, + "query": 769, + "keys": 257, + "heads": 4, + "kv_heads": 1, + "head_dim": 128, + "causal": true + }, + "seed": 1111, + "pattern": "gaussian", + "oracle": { + "dtype": "float32", + "matmul_precision": "highest", + "allow_tf32": false, + "query_chunk": 128 + }, + "rust_kernel": { + "blocks": 16, + "cache": "hot", + "causal": true, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 1, + 769, + 257, + 4, + 1 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 0.0005841255187988281, + "rms": 1.4334200543104076e-05, + "rtol": 1e-05 + }, + "warmup": 10, + "windows": 3 + }, + "validation": { + "sm100_fmha_prefill_cute_rs": { + "max_abs": 0.0005841255187988281, + "rms": 1.433420129615115e-05, + "atol": 0.1, + "rtol": 1e-05, + "elements": 393728 + }, + "sm100_fmha_prefill_cutedsl": { + "max_abs": 0.0005841255187988281, + "rms": 1.4334171282825992e-05, + "atol": 0.1, + "rtol": 1e-05, + "elements": 393728 + }, + "cross_implementation": { + "max_abs": 6.103515625e-05, + "rms": 1.0943016093278857e-07, + "atol": 0.1, + "rtol": 1e-05, + "elements": 393728 + } + }, + "timing": null +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/final-dense.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/final-dense.json new file mode 100644 index 0000000000..41c6e705d4 --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/final-dense.json @@ -0,0 +1,154 @@ +{ + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "source_sha256": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "source_closure_sha256": { + "cute/blackwell/kernel/attention/fmha/fmha.py": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "helpers/fmha_helpers.py": "cbe090eb5567b7cdbd8b4944cbd579084fd60379c463bcfb011e6d4b28ef20e1", + "helpers/__init__.py": "9cf1ae75e6473e5acf38b8ef7b498a14ff74fd54741a930708f6995fcacb7581" + }, + "reference_compiler_optimization_level": 2, + "sanitizer": null, + "rust_executable_sha256": "da252cfe6151cf28e583790078cf8df7125e40e3ae0c88eebf363db76cf19e61", + "config": { + "batch": 4, + "query": 4096, + "keys": 4096, + "heads": 8, + "kv_heads": 8, + "head_dim": 128, + "causal": false + }, + "seed": 1111, + "pattern": "gaussian", + "oracle": { + "dtype": "float32", + "matmul_precision": "highest", + "allow_tf32": false, + "query_chunk": 128 + }, + "rust_kernel": { + "blocks": 148, + "cache": "hot", + "causal": false, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 4, + 4096, + 4096, + 8, + 8 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 1.9516795873641968e-05, + "rms": 2.3820080032120917e-06, + "rtol": 1e-05 + }, + "warmup": 10, + "windows": 3 + }, + "validation": { + "sm100_fmha_prefill_cute_rs": { + "max_abs": 1.9516795873641968e-05, + "rms": 2.3820080059522297e-06, + "atol": 0.1, + "rtol": 1e-05, + "elements": 16777216 + }, + "sm100_fmha_prefill_cutedsl": { + "max_abs": 1.9516795873641968e-05, + "rms": 2.3820159640308702e-06, + "atol": 0.1, + "rtol": 1e-05, + "elements": 16777216 + }, + "cross_implementation": { + "max_abs": 1.52587890625e-05, + "rms": 2.3589489472897185e-08, + "atol": 0.1, + "rtol": 1e-05, + "elements": 16777216 + } + }, + "timing": { + "mode": "direct-cuda-events", + "cache": "hot", + "measurement_scope": "complete invocation throughput including host submission gaps", + "warmup": 2000, + "iterations": 301, + "windows": 5, + "sm100_fmha_prefill_cutedsl": { + "samples_us": [ + 288.99647151909, + 290.04681862865965, + 290.68363227717504, + 290.84428045836796, + 293.76126919869967 + ], + "median_us": 290.68363227717504 + }, + "sm100_fmha_prefill_cute_rs": { + "blocks": 148, + "cache": "hot", + "causal": false, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 301, + "median_us": 290.4759150406847, + "registers_per_thread": 128, + "samples_us": [ + 291.3377831544591, + 290.1817397817821, + 290.4759150406847, + 290.6824663231935, + 290.44889532450424 + ], + "shape": [ + 4, + 4096, + 4096, + 8, + 8 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 1.9516795873641968e-05, + "rms": 2.3820080032120917e-06, + "rtol": 1e-05 + }, + "warmup": 2000, + "windows": 5 + }, + "gpu_before": { + "query": "--query-gpu=name,uuid,driver_version,clocks.current.sm,clocks.current.memory,power.draw,power.limit,utilization.gpu", + "output": "NVIDIA B200, GPU-a4674550-4882-a08d-60b1-5efde3790e39, 615.71.09, 1830, 4000, 195.95, 1000.00, 0", + "error": "" + }, + "gpu_after": { + "query": "--query-gpu=name,uuid,driver_version,clocks.current.sm,clocks.current.memory,power.draw,power.limit,utilization.gpu", + "output": "NVIDIA B200, GPU-a4674550-4882-a08d-60b1-5efde3790e39, 615.71.09, 1830, 4000, 549.16, 1000.00, 0", + "error": "" + } + } +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/memcheck.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/memcheck.json new file mode 100644 index 0000000000..e7e98bbb7d --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/memcheck.json @@ -0,0 +1,87 @@ +{ + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "source_sha256": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "source_closure_sha256": { + "cute/blackwell/kernel/attention/fmha/fmha.py": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "helpers/fmha_helpers.py": "cbe090eb5567b7cdbd8b4944cbd579084fd60379c463bcfb011e6d4b28ef20e1", + "helpers/__init__.py": "9cf1ae75e6473e5acf38b8ef7b498a14ff74fd54741a930708f6995fcacb7581" + }, + "reference_compiler_optimization_level": 2, + "sanitizer": "memcheck", + "rust_executable_sha256": "da252cfe6151cf28e583790078cf8df7125e40e3ae0c88eebf363db76cf19e61", + "config": { + "batch": 1, + "query": 257, + "keys": 385, + "heads": 4, + "kv_heads": 2, + "head_dim": 128, + "causal": true + }, + "seed": 1111, + "pattern": "gaussian", + "oracle": { + "dtype": "float32", + "matmul_precision": "highest", + "allow_tf32": false, + "query_chunk": 128 + }, + "rust_kernel": { + "blocks": 8, + "cache": "hot", + "causal": true, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 1, + 257, + 385, + 4, + 2 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 0.0003203749656677246, + "rms": 1.9929519438751464e-05, + "rtol": 1e-05 + }, + "warmup": 10, + "windows": 3 + }, + "validation": { + "sm100_fmha_prefill_cute_rs": { + "max_abs": 0.0003203749656677246, + "rms": 1.9929519112338312e-05, + "atol": 0.1, + "rtol": 1e-05, + "elements": 131584 + }, + "sm100_fmha_prefill_cutedsl": { + "max_abs": 0.0003203749656677246, + "rms": 1.9929484551539645e-05, + "atol": 0.1, + "rtol": 1e-05, + "elements": 131584 + }, + "cross_implementation": { + "max_abs": 1.52587890625e-05, + "rms": 6.651023909398646e-08, + "atol": 0.1, + "rtol": 1e-05, + "elements": 131584 + } + }, + "timing": null +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/memcheck.log b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/memcheck.log new file mode 100644 index 0000000000..44e00e51ce --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/memcheck.log @@ -0,0 +1,40 @@ +Checking both-tails-gqa: 1,257,385,4,2 +========= COMPUTE-SANITIZER +{ + "blocks": 8, + "cache": "hot", + "causal": true, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 1, + 257, + 385, + 4, + 2 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 0.0003203749656677246, + "rms": 0.000019929519438751464, + "rtol": 0.00001 + }, + "warmup": 10, + "windows": 3 +} +========= ERROR SUMMARY: 0 errors +CuTeDSL: PASS +CuTe-RS: PASS +All selected prefill validation cases passed. diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/opcode-counts.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/opcode-counts.json new file mode 100644 index 0000000000..68395c98fa --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/opcode-counts.json @@ -0,0 +1,244 @@ +{ + "before": { + "sm100_fmha_prefill_cute_rs": { + "ATOMS.AND": 2, + "ATOMS.OR": 2, + "BAR.SYNC.DEFER_BLOCKING": 3, + "BPT.TRAP": 2, + "BRA": 161, + "BRA.U": 21, + "BSSY.RECONVERGENT": 12, + "BSYNC.RECONVERGENT": 12, + "CALL.REL.NOINC": 2, + "CS2R": 1, + "DEPBAR.LE": 4, + "ELECT": 19, + "EXIT": 11, + "F2FP.F16.F32.PACK_AB": 80, + "F2I.FTZ.U32.TRUNC.NTZ": 8, + "FADD": 3, + "FADD2": 64, + "FENCE.VIEW.ASYNC.S": 3, + "FENCE.VIEW.ASYNC.T": 5, + "FFMA2": 64, + "FMNMX": 100, + "FMNMX3": 79, + "FMUL": 5, + "FMUL2": 80, + "FSEL": 257, + "FSETP.NEU.AND": 2, + "I2F.U32.RP": 8, + "IADD3": 2, + "IMAD": 3, + "IMAD.HI.U32": 2, + "IMAD.IADD": 1, + "IMAD.MOV": 3, + "IMAD.MOV.U32": 25, + "IMAD.SHL.U32": 5, + "IMAD.U32": 43, + "ISETP.EQ.U32.AND": 1, + "ISETP.GT.U32.AND": 3, + "ISETP.LE.U32.AND": 131, + "ISETP.LT.U32.AND": 128, + "ISETP.NE.AND": 2, + "ISETP.NE.U32.AND": 12, + "ISETP.NE.U32.AND.EX": 2, + "LDC": 3, + "LDCU": 15, + "LDCU.64": 9, + "LDS": 3, + "LDTM.x16": 10, + "LDTM.x2": 3, + "LDTM.x32": 8, + "LEA": 4, + "LEA.HI": 1, + "LOP3.LUT": 27, + "MEMBAR.ALL.CTA": 2, + "MOV": 2, + "MUFU.EX2": 130, + "MUFU.RCP": 10, + "NANOSLEEP": 1, + "NOP": 94, + "PLOP3.LUT": 4, + "R2UR": 29, + "R2UR.BROADCAST": 1, + "REDUX": 2, + "RET.REL.NODEC": 2, + "S2R": 4, + "S2UR": 3, + "SEL": 1, + "SHF.L.U32": 39, + "SHF.R.U32.HI": 16, + "STS": 1, + "STS.128": 4, + "STS.U8": 1, + "STTM.x16": 8, + "STTM.x2": 2, + "STTM.x32": 2, + "SYNCS.ARRIVE.TRANS64": 6, + "SYNCS.ARRIVE.TRANS64.A1T0": 26, + "SYNCS.EXCH.64": 34, + "SYNCS.PHASECHK.TRANS64.TRYWAIT": 91, + "UFLO.U32": 1, + "UIADD3": 363, + "UIADD3.64": 24, + "UIADD3.X": 38, + "UIMAD": 29, + "UIMAD.WIDE.U32": 23, + "UISETP.EQ.U32.AND": 1, + "UISETP.EQ.U32.OR.EX": 1, + "UISETP.GE.U32.AND": 29, + "UISETP.GE.U32.OR": 1, + "UISETP.GT.U32.AND": 4, + "UISETP.LT.U32.AND": 3, + "UISETP.NE.U32.AND": 19, + "UISETP.NE.U32.AND.EX": 3, + "ULEA": 26, + "ULEA.HI": 6, + "ULEA.HI.X": 1, + "ULOP3.LUT": 102, + "UMOV": 129, + "UPLOP3.LUT": 3, + "USEL": 8, + "USETMAXREG.DEALLOC.CTAPOOL": 5, + "USETMAXREG.TRY_ALLOC.CTAPOOL": 1, + "USHF.L.U32": 23, + "USHF.R.U32.HI": 34, + "UTCATOMSWS.AND": 1, + "UTCATOMSWS.FIND_AND_SET.ALIGN": 2, + "UTCBAR": 18, + "UTCHMMA": 64, + "UTMACMDFLUSH": 2, + "UTMALDG.3D": 6, + "UTMASTG.3D": 2, + "UVIRTCOUNT.DEALLOC.SMPOOL": 1, + "VIADD": 33, + "VIADDMNMX.U32": 1, + "VOTE.ALL": 1, + "VOTEU.ALL": 1, + "VOTEU.ANY": 3, + "WARPSYNC.ALL": 1 + } + }, + "after": { + "sm100_fmha_prefill_cute_rs": { + "ATOMS.AND": 2, + "ATOMS.OR": 2, + "BAR.SYNC.DEFER_BLOCKING": 3, + "BPT.TRAP": 2, + "BRA": 161, + "BRA.U": 21, + "BSSY.RECONVERGENT": 12, + "BSYNC.RECONVERGENT": 12, + "CALL.REL.NOINC": 2, + "CS2R": 1, + "DEPBAR.LE": 4, + "ELECT": 19, + "EXIT": 11, + "F2FP.F16.F32.PACK_AB": 80, + "F2I.FTZ.U32.TRUNC.NTZ": 8, + "FADD": 3, + "FADD2": 64, + "FENCE.VIEW.ASYNC.S": 3, + "FENCE.VIEW.ASYNC.T": 5, + "FFMA2": 64, + "FMNMX": 100, + "FMNMX3": 79, + "FMUL": 5, + "FMUL2": 80, + "FSEL": 257, + "FSETP.NEU.AND": 2, + "I2F.U32.RP": 8, + "IADD3": 2, + "IMAD": 3, + "IMAD.HI.U32": 2, + "IMAD.IADD": 1, + "IMAD.MOV": 3, + "IMAD.MOV.U32": 27, + "IMAD.SHL.U32": 5, + "IMAD.U32": 43, + "ISETP.EQ.U32.AND": 1, + "ISETP.GT.U32.AND": 3, + "ISETP.LE.U32.AND": 131, + "ISETP.LT.U32.AND": 128, + "ISETP.NE.AND": 2, + "ISETP.NE.U32.AND": 12, + "ISETP.NE.U32.AND.EX": 2, + "LDC": 3, + "LDCU": 15, + "LDCU.64": 9, + "LDS": 4, + "LDTM.x16": 10, + "LDTM.x2": 3, + "LDTM.x32": 8, + "LEA": 4, + "LEA.HI": 1, + "LOP3.LUT": 28, + "MEMBAR.ALL.CTA": 2, + "MOV": 2, + "MUFU.EX2": 130, + "MUFU.RCP": 10, + "NANOSLEEP": 1, + "NOP": 100, + "PLOP3.LUT": 4, + "R2UR": 30, + "R2UR.BROADCAST": 1, + "REDUX": 2, + "RET.REL.NODEC": 2, + "S2R": 4, + "S2UR": 3, + "SEL": 1, + "SHF.L.U32": 41, + "SHF.R.U32.HI": 17, + "STS": 1, + "STS.128": 4, + "STS.U8": 1, + "STTM.x16": 8, + "STTM.x2": 2, + "STTM.x32": 2, + "SYNCS.ARRIVE.TRANS64": 6, + "SYNCS.ARRIVE.TRANS64.A1T0": 26, + "SYNCS.EXCH.64": 34, + "SYNCS.PHASECHK.TRANS64.TRYWAIT": 91, + "UFLO.U32": 1, + "UIADD3": 363, + "UIADD3.64": 24, + "UIADD3.X": 38, + "UIMAD": 29, + "UIMAD.WIDE.U32": 23, + "UISETP.EQ.U32.AND": 1, + "UISETP.EQ.U32.OR.EX": 1, + "UISETP.GE.U32.AND": 29, + "UISETP.GE.U32.OR": 1, + "UISETP.GT.U32.AND": 4, + "UISETP.LT.U32.AND": 3, + "UISETP.NE.U32.AND": 19, + "UISETP.NE.U32.AND.EX": 3, + "ULEA": 26, + "ULEA.HI": 6, + "ULEA.HI.X": 1, + "ULOP3.LUT": 101, + "UMOV": 127, + "UPLOP3.LUT": 3, + "USEL": 8, + "USETMAXREG.DEALLOC.CTAPOOL": 5, + "USETMAXREG.TRY_ALLOC.CTAPOOL": 1, + "USHF.L.U32": 21, + "USHF.R.U32.HI": 33, + "UTCATOMSWS.AND": 1, + "UTCATOMSWS.FIND_AND_SET.ALIGN": 2, + "UTCBAR": 18, + "UTCHMMA": 64, + "UTMACMDFLUSH": 2, + "UTMALDG.3D": 6, + "UTMASTG.3D": 2, + "UVIRTCOUNT.DEALLOC.SMPOOL": 1, + "VIADD": 33, + "VIADDMNMX.U32": 1, + "VOTE.ALL": 1, + "VOTEU.ALL": 1, + "VOTEU.ANY": 3, + "WARPSYNC.ALL": 1 + } + } +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/summary.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/summary.json new file mode 100644 index 0000000000..6c48435970 --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/summary.json @@ -0,0 +1,341 @@ +{ + "schema_version": 1, + "recorded_date_utc": "2026-10-14", + "example": "sm100_fmha_prefill_cute", + "baseline_commit": "d642dcbfbf8c2f249a6f07307c158c5c0e89a4f7", + "source_provenance": { + "description": "Working-tree changed source and fixture hashes relative to baseline; artifact and Markdown files excluded.", + "sha256": { + "cuda-oxide/cuteir/ctm-rs/src/async_copy.rs": "7216d9b0416cc9bdc63f5b8e9bf72e257e6114ebf7df1345bcf9050db8f227b0", + "cuda-oxide/cuteir/ctm-rs/src/f32x2.rs": "b277bc8a515154f36f306cc0f99336bbcf2438c3b8e46389ef8b771cc08c6669", + "cuda-oxide/cuteir/ctm-rs/src/float.rs": "85baddeb2ca6fe8335d85f3d362b5732cbbea42843e87d2e81028a88312ae31c", + "cuda-oxide/cuteir/ctm-rs/src/lib.rs": "1484938ba4ae30650586063ff759245cf59e4a5acc4a99b2240a9d096a30c4b0", + "cuda-oxide/cuteir/ctm-rs/src/mbarrier.rs": "8cd83101598d2ae2c07b2c011b1f304c682ea9d6d41f5ae6c979277c138d1071", + "cuda-oxide/cuteir/ctm-rs/src/memory.rs": "9a57659b2267db145f44554d1f579b7e41ab46fc9d1fc1743efa425a4644e5a6", + "cuda-oxide/cuteir/ctm-rs/src/tcgen05/mod.rs": "c8e7359f228942f15e9e434039351c9989d3641984b0533c4cb9bc14e1bb532c", + "cuda-oxide/cuteir/ctm-rs/src/tcgen05/tmem.rs": "4e2f906185232d8c67080bd368ad3cd59d154f340a221807099c7c5d178c07b7", + "cuda-oxide/cuteir/ctm-rs/src/warp.rs": "47cbf3c4688b5b60608749b9fef86003db03eccc8d64fa78f0d8c820e3c91741", + "cuda-oxide/cuteir/ctm-rs/tests/attention_api.rs": "56d0c52eeb8515d18df0ae8100a596c7dea3dd7d9d3b8c73cab421e4d5a331f1", + "cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_alignment.rs": "fa3534ac9c4e4b118116e95fc3e7f1d24d25cdd0aa1729847a3f02ca836e8d1f", + "cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_alignment.stderr": "58e7ee175482a544009c165c14a35462a3c7501e12118db32d6f2741748737b6", + "cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_range.rs": "f5a372f0642b530e73d694f0cb6afef7ac3703aef72272251d8fc03649e8c482", + "cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_range.stderr": "6517c96a4904c3d44c322e76ca2413e10ce3c54510a46fefbc0989fa85d64cfa", + "cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_runtime.rs": "471a632a24bb66ca39bce88beb32a8d02e62bbc1f76207f3a8fb16a8317b8421", + "cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_register_count_runtime.stderr": "601a08b9a76696f70a3def2e57c876ea8cf97434f8f7bd089235a43f33fe1ca9", + "cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.rs": "efb2faae76a483784499fcca55e73de99ed803a94933a944d7f68d34f022edca", + "cuda-oxide/cuteir/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.stderr": "d1e50904ff7df112feba6f08baa08d3ed8143d436825a96626b9e0cdfbdbdaac", + "cuda-oxide/cuteir/ctm-rs/tests/compile_pass/sm100_attention_primitives.rs": "5f4878047736bdd81daf46695b9b607ab4b8ba6042981a92cb99183782b1624a", + "cuda-oxide/cuteir/cute-rs/src/arch.rs": "3a31898e23e1614e930f7e2b139740bccf5ab0eaded23fbef6dbe3f21e756d0c", + "cuda-oxide/cuteir/cute-rs/src/sm100_attention/__compiler.rs": "568b86f8d490dc02831de1fdb96ec033b2c6016f1891290d190feb10175462d6", + "cuda-oxide/cuteir/cute-rs/src/sm100_attention/mod.rs": "25226967dff2835d3c4e92c32c7f59e47f18ef0ac5b896e06669721875bed2bc", + "cuda-oxide/cuteir/cute-rs/src/sm100_attention/pipeline.rs": "731fc189033505e140bffdd5d40c6ce7013ed5966c2c36347b37c943873396e0", + "cuda-oxide/cuteir/cute-rs/src/sm100_attention/tma.rs": "249955aa56a9d503909d08abe3543cb938641c209924aeb93bcb6176c361ca2f", + "cuda-oxide/cuteir/cute-rs/src/sm100_attention/tmem.rs": "e0f0d35cf72f29d1ad101c034a7961c66daca2af16ecc2d1fad154e1e528b02f", + "cuda-oxide/cuteir/cute-rs/tests/sm100_attention.rs": "e0560e78198b0bf0d156c1013d13048d045c9e8b97e0b0c6fe2c548d9293348b", + "cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/src/kernel.rs": "5dd43415a4735fd904d9570ebc58d76cb1e7a69a0275ad5bb7df796fcc4793e3" + }, + "final_rebuild_confirmation": "Final rebuilt cubin is byte-identical to the migrated cubin used for the correctness matrix.", + "final_binary_sha256": "da252cfe6151cf28e583790078cf8df7125e40e3ae0c88eebf363db76cf19e61", + "final_cubin_sha256": "3229187c799d4693e290e252e56ed086bb5e8ae67e0e66f8693269e3509404ce" + }, + "environment": { + "python": "3.12.3 (main, Jul 15 2026, 23:46:41) [GCC 13.3.0]", + "torch": "2.13.0+cu130", + "cutlass": "4.7.0", + "compiler": "CUTLASS 4.7.0", + "architecture": "sm_100a", + "reference_optimization_level": 2, + "rust_optimization_level": "2", + "cuda_visible_devices": null, + "cpu_affinity": [ + 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 + ], + "gpu": "NVIDIA B200", + "driver": "615.71.09", + "host_cuda_toolkit": "13.4", + "rust_toolchain": "nightly-2026-10-11", + "device_backend": "cutlass-mlir", + "clocks_locked": false + }, + "reference": { + "commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "sha256": { + "cute/blackwell/kernel/attention/fmha/fmha.py": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "helpers/fmha_helpers.py": "cbe090eb5567b7cdbd8b4944cbd579084fd60379c463bcfb011e6d4b28ef20e1", + "helpers/__init__.py": "9cf1ae75e6473e5acf38b8ef7b498a14ff74fd54741a930708f6995fcacb7581" + } + }, + "correctness": { + "status": "passed", + "passed_cases": 8, + "raw_reports": [ + "correctness/both-tails-gqa.json", + "correctness/kv-ring-wrap.json", + "correctness/persistent-reuse.json", + "correctness/query-tail.json", + "correctness/reference-smoke.json", + "correctness/short-causal.json", + "correctness/single-key.json", + "correctness/unequal-causal.json" + ], + "worst_case_errors": { + "max_abs_rust_oracle": 0.0005841255187988281, + "max_abs_cross_implementation": 6.103515625e-05 + }, + "cases": [ + { + "name": "both-tails-gqa", + "configuration": { + "batch": 1, + "query": 257, + "keys": 385, + "heads": 4, + "kv_heads": 2, + "head_dim": 128, + "causal": true + }, + "seed": 1111, + "pattern": "gaussian" + }, + { + "name": "kv-ring-wrap", + "configuration": { + "batch": 1, + "query": 512, + "keys": 1153, + "heads": 8, + "kv_heads": 2, + "head_dim": 128, + "causal": false + }, + "seed": 1111, + "pattern": "gaussian" + }, + { + "name": "persistent-reuse", + "configuration": { + "batch": 4, + "query": 4096, + "keys": 4096, + "heads": 8, + "kv_heads": 8, + "head_dim": 128, + "causal": false + }, + "seed": 1111, + "pattern": "gaussian" + }, + { + "name": "query-tail", + "configuration": { + "batch": 1, + "query": 129, + "keys": 128, + "heads": 2, + "kv_heads": 2, + "head_dim": 128, + "causal": false + }, + "seed": 1111, + "pattern": "gaussian" + }, + { + "name": "reference-smoke", + "configuration": { + "batch": 1, + "query": 1024, + "keys": 1024, + "heads": 8, + "kv_heads": 8, + "head_dim": 128, + "causal": false + }, + "seed": 1111, + "pattern": "gaussian" + }, + { + "name": "short-causal", + "configuration": { + "batch": 1, + "query": 17, + "keys": 31, + "heads": 4, + "kv_heads": 2, + "head_dim": 128, + "causal": true + }, + "seed": 1111, + "pattern": "uniform" + }, + { + "name": "single-key", + "configuration": { + "batch": 1, + "query": 1, + "keys": 1, + "heads": 1, + "kv_heads": 1, + "head_dim": 128, + "causal": false + }, + "seed": 1111, + "pattern": "uniform" + }, + { + "name": "unequal-causal", + "configuration": { + "batch": 1, + "query": 769, + "keys": 257, + "heads": 4, + "kv_heads": 1, + "head_dim": 128, + "causal": true + }, + "seed": 1111, + "pattern": "gaussian" + } + ] + }, + "machine_code_comparison": { + "status": "passed_with_teardown_difference", + "kernel_text_byte_identical": false, + "added_non_padding_instructions": 2, + "reason": "Typed TMEM deallocation reloads the shared allocation token; the change is outside the attention iteration loops.", + "resources_unchanged": true, + "registers": 128, + "stack_bytes": 0, + "local_bytes": 0, + "static_shared_bytes": 1024, + "kernel_text_sha256_before": "9d9380131b0fbf65b133ac48e7fcb7c5c5125dc66bfc7884c6d155d40bbe1ff6", + "kernel_text_sha256_after": "a130023d7ac17228ec32b9cb2f40bbaaeaa516b64fddbced84237026a4f310e7" + }, + "sanitizers": { + "status": "passed", + "runs": [ + { + "tool": "memcheck", + "errors": 0, + "configuration": { + "batch": 1, + "query": 257, + "keys": 385, + "heads": 4, + "kv_heads": 2, + "head_dim": 128, + "causal": true + }, + "raw_report": "memcheck.json" + }, + { + "tool": "synccheck", + "errors": 0, + "configuration": { + "batch": 1, + "query": 512, + "keys": 1153, + "heads": 8, + "kv_heads": 2, + "head_dim": 128, + "causal": false + }, + "raw_report": "synccheck.json" + } + ], + "raw_reports": [ + "memcheck.log", + "memcheck.json", + "synccheck.log", + "synccheck.json" + ] + }, + "timing": { + "status": "completed_one_profile", + "configuration": { + "batch": 4, + "query": 4096, + "keys": 4096, + "heads": 8, + "kv_heads": 8, + "head_dim": 128, + "causal": false + }, + "method": "direct-cuda-events", + "measurement_scope": "complete invocation throughput including host submission gaps", + "cache": "hot", + "warmup": 2000, + "iterations": 301, + "windows": 5, + "baseline": { + "median_us": 290.3222879302066, + "range_us": [ + 286.4804695610984, + 290.72340144667515 + ] + }, + "final": { + "median_us": 290.4759150406847, + "range_us": [ + 290.1817397817821, + 291.3377831544591 + ] + }, + "relative_change_percent": 0.05291605807233868, + "raw_reports": [ + "baseline-dense.json", + "final-dense.json" + ], + "conclusion": "No measurable regression in this run; the 0.053% median increase is within the observed spread.", + "parity_claim": false, + "limits": "One dense persistent profile; complete-invocation throughput includes host submission gaps. This does not establish broad performance parity." + }, + "evidence_limits": "Correctness covers the recorded configurations. Static instruction comparison does not establish timing parity.", + "artifact_sha256": { + "correctness/both-tails-gqa.json": "f62136d05b59aedc3b699a9f0ac36884cd7128f4f8ff05f6b16983c2aef2a0b8", + "correctness/kv-ring-wrap.json": "3d29a80a31151ed3320b3f8682e9e5764cc61c14f3890c6448ac12f5b96e02e8", + "correctness/persistent-reuse.json": "44efc9dcae67073794c4ab8021c3bc1beaa9d60a22ceb47a270b5b849c536ee8", + "correctness/query-tail.json": "af283eeb8b39366c39499f7eee9cfb60dd9e9c2c1bc9828fae0ca61afd13e178", + "correctness/reference-smoke.json": "628367894637ee9f8abc63241e685692b0b68adb5912faf7eb8c040b71cbbd9d", + "correctness/short-causal.json": "a6709460ec4234bb3f3fd768dc14234fdb9ba4185ae67f2a90e3be4dc5e88632", + "correctness/single-key.json": "eb0a431b88222f0f2205536e50e75f70c2985b12b800b3c337e6b22a39afa989", + "correctness/unequal-causal.json": "fb58ec0ff2e672d7d5c80a8d7167ef71716bd6094397953357e68d4cc28c8fa7", + "opcode-counts.json": "ba7d2bed70eec0e4cc5321d13f5b9d864c05e313191543bf22a7642333836e83", + "memcheck.log": "20dcdda748bff5bb28163d53ac9c26d36e0c0341fb56d0e8f4be6b903e6c919f", + "memcheck.json": "6557857dea98d46af4dca140e800054ddd10086520e2fe96a216353cf4900726", + "synccheck.log": "6edc23ea3718e7f0dbecadabcad0d5a25ff40d62899f84d3eb83779fa9a620bf", + "synccheck.json": "10b35a55b2ec4b282c06e59e829b9d2353fa87a5c9c3501e17609654266ecb05", + "baseline-dense.json": "0f2a56332a72f5c90ecc76c2f82a1f48049043b08dddf8d157e1cb15a87e1b08", + "final-dense.json": "ed4cffd8cc973f6bd2fe81c8ffa8117c9aae64a3b2d4cdea328132efb11f176d" + } +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/synccheck.json b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/synccheck.json new file mode 100644 index 0000000000..b869673f34 --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/synccheck.json @@ -0,0 +1,87 @@ +{ + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "source_sha256": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "source_closure_sha256": { + "cute/blackwell/kernel/attention/fmha/fmha.py": "c11801367b5681e531b756238920124ba4d2ca7809bdb4f9aacd9aaa9ae3a955", + "helpers/fmha_helpers.py": "cbe090eb5567b7cdbd8b4944cbd579084fd60379c463bcfb011e6d4b28ef20e1", + "helpers/__init__.py": "9cf1ae75e6473e5acf38b8ef7b498a14ff74fd54741a930708f6995fcacb7581" + }, + "reference_compiler_optimization_level": 2, + "sanitizer": "synccheck", + "rust_executable_sha256": "da252cfe6151cf28e583790078cf8df7125e40e3ae0c88eebf363db76cf19e61", + "config": { + "batch": 1, + "query": 512, + "keys": 1153, + "heads": 8, + "kv_heads": 2, + "head_dim": 128, + "causal": false + }, + "seed": 1111, + "pattern": "gaussian", + "oracle": { + "dtype": "float32", + "matmul_precision": "highest", + "allow_tf32": false, + "query_chunk": 128 + }, + "rust_kernel": { + "blocks": 16, + "cache": "hot", + "causal": false, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 1, + 512, + 1153, + 8, + 2 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 2.6892870664596558e-05, + "rms": 4.552741129992708e-06, + "rtol": 1e-05 + }, + "warmup": 10, + "windows": 3 + }, + "validation": { + "sm100_fmha_prefill_cute_rs": { + "max_abs": 2.6892870664596558e-05, + "rms": 4.55274130217731e-06, + "atol": 0.1, + "rtol": 1e-05, + "elements": 524288 + }, + "sm100_fmha_prefill_cutedsl": { + "max_abs": 2.6892870664596558e-05, + "rms": 4.552762220555451e-06, + "atol": 0.1, + "rtol": 1e-05, + "elements": 524288 + }, + "cross_implementation": { + "max_abs": 1.52587890625e-05, + "rms": 2.7120522716472806e-08, + "atol": 0.1, + "rtol": 1e-05, + "elements": 524288 + } + }, + "timing": null +} diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/synccheck.log b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/synccheck.log new file mode 100644 index 0000000000..ae1b61a04b --- /dev/null +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/reference/results/ctm-migration/synccheck.log @@ -0,0 +1,40 @@ +Checking kv-ring-wrap: 1,512,1153,8,2 +========= COMPUTE-SANITIZER +{ + "blocks": 16, + "cache": "hot", + "causal": false, + "compiler": "CUTLASS 4.7.0", + "compiler_optimization_level": 2, + "gpu": "NVIDIA B200", + "head_dim": 128, + "implementation": "sm100_fmha_prefill_cute_rs", + "iterations": 100, + "median_us": null, + "registers_per_thread": 128, + "samples_us": [], + "shape": [ + 1, + 512, + 1153, + 8, + 2 + ], + "shared_memory_bytes": 230400, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53", + "target": "sm_100a", + "threads": 512, + "timing": "direct-cuda-events", + "validation": { + "atol": 0.10000000149011612, + "max_abs": 0.000026892870664596558, + "rms": 4.552741129992708e-6, + "rtol": 0.00001 + }, + "warmup": 10, + "windows": 3 +} +========= ERROR SUMMARY: 0 errors +CuTeDSL: PASS +CuTe-RS: PASS +All selected prefill validation cases passed. diff --git a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/src/kernel.rs b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/src/kernel.rs index 58b479bc7b..3aa3d3db1d 100644 --- a/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/src/kernel.rs +++ b/cuda-oxide/cuteir/examples/sm100_fmha_prefill_cute/src/kernel.rs @@ -7,14 +7,17 @@ //! Two independent softmax groups exchange the aliased S/P tensor-memory //! regions. A separate correction group rescales the running output before //! the MMA warp accumulates the next FP16 probability tile. +//! CuTe supplies typed MMA/TMA operations; CTM supplies explicit instructions, +//! shared storage, and the kernel launch macros. -use ctm_rs::{thread, warp}; -use cuda_device::barrier::Barrier; -use cuda_device::f32x2::{ +use ctm_rs::f32x2::{ add_f32x2, fma_f32x2, mul_f32x2, pack_f32x2, unpack_f32x2_hi, unpack_f32x2_lo, }; -use cuda_device::{ - DynamicSharedArray, barrier, cuda_module, kernel, launch_bounds, launch_contract, tcgen05, +use ctm_rs::mbarrier::{self as barrier, Barrier}; +use ctm_rs::memory::{DynamicSharedArray, SharedAddress, U32x4}; +use ctm_rs::tcgen05::tmem::{self, TmemGuard, TmemReady, TmemUninit}; +use ctm_rs::{ + convert, cuda_module, float, kernel, launch_bounds, launch_contract, tcgen05, thread, warp, }; use cute_rs::sm100_attention::{MmaSs, MmaTs, SharedLayout, SharedTile, TmaCopy, TmemFragment}; use cute_rs::{OperandA, OperandB, TmaDesc}; @@ -94,7 +97,7 @@ impl Ring { #[inline(always)] unsafe fn wait(self, count: u32) { unsafe { wait(self.full_ptr(count), (count / STAGES) & 1) }; - tcgen05::tcgen05_fence_after_thread_sync(); + tcgen05::fence_after_thread_sync(); } /// Select the full barrier for the monotonically increasing stage count. #[inline(always)] @@ -104,7 +107,7 @@ impl Ring { /// Order prior TMEM accesses and contribute one thread arrival to the full barrier. #[inline(always)] unsafe fn commit_threads(self, count: u32) { - tcgen05::tcgen05_fence_before_thread_sync(); + tcgen05::fence_before_thread_sync(); unsafe { barrier::mbarrier_arrive(self.full_ptr(count)); } @@ -112,7 +115,7 @@ impl Ring { /// Order prior TMEM accesses and contribute one thread arrival to the empty barrier. #[inline(always)] unsafe fn release_threads(self, count: u32) { - tcgen05::tcgen05_fence_before_thread_sync(); + tcgen05::fence_before_thread_sync(); unsafe { barrier::mbarrier_arrive(self.empty.add((count % STAGES) as usize)); } @@ -123,7 +126,7 @@ impl Ring { unsafe fn commit_mma(self, count: u32) { if warp::is_elected_sync(u32::MAX) { unsafe { - tcgen05::tcgen05_commit(self.full_ptr(count).cast()); + tcgen05::commit(SharedAddress::from_raw(self.full_ptr(count))); } } } @@ -133,7 +136,9 @@ impl Ring { unsafe fn release_mma(self, count: u32) { if warp::is_elected_sync(u32::MAX) { unsafe { - tcgen05::tcgen05_commit(self.empty.add((count % STAGES) as usize).cast()); + tcgen05::commit(SharedAddress::from_raw( + self.empty.add((count % STAGES) as usize), + )); } } } @@ -246,7 +251,7 @@ fn tmem_rows(base: u32, tid: u32) -> u32 { /// Evaluate the device approximate base-two exponential with subnormal flushing. #[inline(always)] fn exp2(x: f32) -> f32 { - cuda_device::float::ex2_approx_ftz_f32(x) + float::ex2_approx_ftz_f32(x) } /// Source SM100 max reduction has two independent chains per 32-score chunk. @@ -254,7 +259,7 @@ fn exp2(x: f32) -> f32 { fn reduce_score_chunk(scores: &TmemFragment<32>, mut left: f32, mut right: f32) -> (f32, f32) { let mut i = 0; while i < 32 { - cuda_device::thread::__unroll_config::<0>(); + thread::__unroll_config::<0>(); left = left.max(scores.values[i]); left = left.max(scores.values[i + 1]); right = right.max(scores.values[i + 2]); @@ -270,7 +275,7 @@ fn reduce_masked_scores(scores: &TmemFragment<128>) -> f32 { let mut maxima = [f32::NEG_INFINITY; 64]; let mut i = 0; while i < 64 { - cuda_device::thread::__unroll_config::<0>(); + thread::__unroll_config::<0>(); maxima[i] = scores.values[i * 2].max(scores.values[i * 2 + 1]); i += 1; } @@ -288,7 +293,7 @@ fn reduce_masked_scores(scores: &TmemFragment<128>) -> f32 { fn reduce_max_level(maxima: &mut [f32; 64]) { let mut i = 0; while i < WIDTH { - cuda_device::thread::__unroll_config::<0>(); + thread::__unroll_config::<0>(); maxima[i] = maxima[i * 2].max(maxima[i * 2 + 1]); i += 1; } @@ -298,7 +303,7 @@ fn reduce_max_level(maxima: &mut [f32; 64]) { #[inline(always)] fn scaled_half_pair(lo: f32, hi: f32, factor: u64) -> u32 { let scaled = mul_f32x2(pack_f32x2(lo, hi), factor); - cuda_device::convert::cvt_f16x2_f32(unpack_f32x2_lo(scaled), unpack_f32x2_hi(scaled)) + convert::cvt_f16x2_f32(unpack_f32x2_lo(scaled), unpack_f32x2_hi(scaled)) } /// Acquire score/statistics slots and issue QK without accumulation, then publish score completion. @@ -398,7 +403,7 @@ pub mod kernels { barrier::fence_mbarrier_init_release_cluster(); barrier::barrier_cta_sync(0, 512); if warp_id == 8 { - tcgen05::tcgen05_alloc(token, 512); + TmemGuard::::from_raw_slot(token).alloc_collective(); } let mut tmem = 0; if warp_id <= 12 { @@ -407,7 +412,7 @@ pub mod kernels { } if warp_id == 13 { - cuda_device::thread::setmaxnreg_dec::<32>(); + thread::setmaxnreg_dec::<32>(); let mut qcount = 0u32; let mut kvcount = 0u32; while tile < total { @@ -516,7 +521,7 @@ pub mod kernels { tile += stride; } } else if warp_id == 12 { - cuda_device::thread::setmaxnreg_dec::<32>(); + thread::setmaxnreg_dec::<32>(); let (mut qcount, mut kvcount, mut scount, mut pcount, mut stats_count, mut ocount) = (0u32, 0u32, 0u32, 0u32, 0u32, 0u32); while tile < total { @@ -644,7 +649,7 @@ pub mod kernels { tile += stride; } } else if warp_id < 8 { - cuda_device::thread::setmaxnreg_inc::<192>(); + thread::setmaxnreg_inc::<192>(); let stage = warp_id / 4; let (scores, probabilities, stats, alias_producer, alias_consumer) = if stage == 0 { (pipes.s0, pipes.p0, pipes.stats0, pipes.alias0, pipes.alias1) @@ -704,7 +709,7 @@ pub mod kernels { let (left, right) = reduce_score_chunk(&third, left, right); // Match the source's handoff before the last max // chunk, permitting the aliased P writer to advance. - tcgen05::tcgen05_store_wait(); + tmem::store_wait(); scores.release_threads(count); alias_producer.commit_threads(alias_send); let (left, right) = reduce_score_chunk(&fourth, left, right); @@ -760,7 +765,7 @@ pub mod kernels { if col >= 2 { let x = values.values[(col - 2) * 2]; let y = values.values[(col - 2) * 2 + 1]; - packed[col - 2] = cuda_device::convert::cvt_f16x2_f32(x, y); + packed[col - 2] = convert::cvt_f16x2_f32(x, y); sums = add_f32x2(sums, pack_f32x2(x, y)); } values.values[col * 2] = exp2(values.values[col * 2]); @@ -782,7 +787,7 @@ pub mod kernels { while col < 64 { let x = values.values[col * 2]; let y = values.values[col * 2 + 1]; - packed[col] = cuda_device::convert::cvt_f16x2_f32(x, y); + packed[col] = convert::cvt_f16x2_f32(x, y); sums = add_f32x2(sums, pack_f32x2(x, y)); col += 1; } @@ -810,7 +815,7 @@ pub mod kernels { tile += stride; } } else if warp_id < 12 { - cuda_device::thread::setmaxnreg_dec::<96>(); + thread::setmaxnreg_dec::<96>(); let (mut stat_count, mut ocount, mut epicount) = (0u32, 0u32, 0u32); while tile < total { let row = (tile % qtiles) * 256; @@ -834,7 +839,7 @@ pub mod kernels { TmemFragment::<2>::load(tmem_rows(tmem + stage * 128, tid)); stats.release_threads(stat_count); pipes.output.wait(ocount); - if !cuda_device::warp::all(maxima.values[0] == maxima.values[1]) { + if !warp::all(maxima.values[0] == maxima.values[1]) { let factor = exp2((maxima.values[0] - maxima.values[1]) * SCALE_LOG2); let factor = pack_f32x2(factor, factor); @@ -860,7 +865,7 @@ pub mod kernels { out.store_deferred(addr); col += 16; } - tcgen05::tcgen05_store_wait(); + tmem::store_wait(); } pipes.output.release_threads(ocount); ocount += 1; @@ -881,7 +886,7 @@ pub mod kernels { stats.release_threads(stat_count); pipes.output.wait(ocount); pipes.epi.acquire(epicount); - let factor = cuda_device::float::rcp_approx_ftz_f32(sums.values[0]); + let factor = float::rcp_approx_ftz_f32(sums.values[0]); let factor = pack_f32x2(factor, factor); let dst = smem .add(O_OFFSET + (epicount % 2) as usize * TILE_BYTES) @@ -899,14 +904,14 @@ pub mod kernels { let mut j = 0usize; #[unroll] while j < 16 { - let words = cuda_device::vector::U32x4([ + let words = U32x4([ scaled_half_pair(out.values[j], out.values[j + 1], factor), scaled_half_pair(out.values[j + 2], out.values[j + 3], factor), scaled_half_pair(out.values[j + 4], out.values[j + 5], factor), scaled_half_pair(out.values[j + 6], out.values[j + 7], factor), ]); *dst.add(OLayout::offset(tid % 128, col + j as u32)) - .cast::() = words; + .cast::() = words; j += 8; } col += 16; @@ -923,11 +928,11 @@ pub mod kernels { } barrier::barrier_cta_sync(7, 128); if warp_id == 8 { - tcgen05::tcgen05_relinquish_alloc_permit(); - tcgen05::tcgen05_dealloc(tmem, 512); + tcgen05::relinquish_alloc_permit(); + TmemGuard::::adopt_ready(token).dealloc_collective(); } } else if warp_id == 14 { - cuda_device::thread::setmaxnreg_dec::<32>(); + thread::setmaxnreg_dec::<32>(); let mut count = 0u32; while tile < total { let row = (tile % qtiles) * 256; @@ -948,18 +953,18 @@ pub mod kernels { ctm_rs::tma::bulk_commit_group(); ctm_rs::tma::bulk_wait_group_read::<1>(); } - cuda_device::warp::sync_mask(u32::MAX); + warp::sync_mask(u32::MAX); pipes.epi.release_threads(count); if warp::is_elected_sync(u32::MAX) { ctm_rs::tma::bulk_wait_group_read::<0>(); } - cuda_device::warp::sync_mask(u32::MAX); + warp::sync_mask(u32::MAX); pipes.epi.release_threads(count + 1); count += 2; tile += stride; } } else { - cuda_device::thread::setmaxnreg_dec::<32>(); + thread::setmaxnreg_dec::<32>(); } } } diff --git a/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/README.md b/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/README.md index 7ed97752fe..6f3bf4ea8c 100644 --- a/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/README.md +++ b/cuda-oxide/cuteir/examples/sm100_mla_decode_fp16_cute/README.md @@ -4,6 +4,14 @@ This cuda-oxide example ports the public CUTLASS FP16 MLA decode kernel to CuTe-RS for data center Blackwell. Two-CTA clusters process a paged latent KV cache; a second kernel combines partial results when multiple splits are used. +The kernel uses `cute_rs` for typed layouts, structured MMA and TMA operations, +and `ctm_rs` for explicit instructions: pipeline synchronization, TMEM transfers, +packed arithmetic, warp/cluster coordination and register control. CTM also +exposes the kernel macros and canonical storage types. The manifest retains +`cuda-device` because the macro expansions still refer to that crate internally. +CuTe operations retain their CuTe dialect lowering through `cutlass-mlir`; +this example uses both APIs rather than the primitives-only compiler profile. + - **128 heads**, one query, latent dimension **512**, RoPE dimension **64**, and **128-token pages**; the latent cache supplies both K and V. - FP16 inputs, probabilities and output; FP32 QK/PV accumulation and base-2 LSE. @@ -14,6 +22,12 @@ cache; a second kernel combines partial results when multiple splits are used. From `cuda-oxide/`, with CUTLASS 4.7 installed (`cargo oxide toolchain install cutlass`): ```bash +# Compile without launching a GPU kernel +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 cargo oxide build sm100_mla_decode_fp16_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_mla_decode_fp16_cute --arch sm_100a -- --batch 1 --kv-length 1024 --splits 1 --clusters 1 --check-only + # Small, unsplit decode CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 cargo oxide run sm100_mla_decode_fp16_cute --arch sm_100a -- --batch 1 --kv-length 1024 --splits 1 --clusters 1 --warmup 300 --iters 100 --windows 5 @@ -33,8 +47,8 @@ CuTe-RS:
+

pliron-mlir-interop

+

IR translation between pliron and MLIR.

+

+ Redirecting to the + crate documentation… +

+
-

pliron-mlir-interop

-

IR translation between pliron and MLIR.

-

- Redirecting to the - crate documentation… -

-