From 5620feea4815a5d74077b836ae4d083ef940572e Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Sun, 6 Sep 2026 07:59:15 -0700 Subject: [PATCH] fix(tma): preserve setmaxnreg with local G2S Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- Cargo.lock | 1 + crates/cuda-oxide-codegen/src/iket.rs | 1 + crates/cuda-oxide-codegen/src/lower.rs | 2 + crates/cuda-oxide-codegen/src/pipeline.rs | 57 +- crates/mir-lower/Cargo.toml | 1 + crates/mir-lower/src/context.rs | 2 +- .../mir-lower/src/convert/intrinsics/tma.rs | 89 ++- .../mir-lower/src/convert/ops/arithmetic.rs | 1 + crates/mir-lower/src/convert/ops/call.rs | 3 + crates/mir-lower/src/lib.rs | 6 +- crates/mir-lower/tests/lowering_test/main.rs | 1 + crates/mir-lower/tests/lowering_test/tma.rs | 145 ++++ crates/rustc-codegen-cuda/Cargo.lock | 1 + .../examples/tma_setmaxnreg_repro/Cargo.lock | 733 ++++++++++++++++++ .../examples/tma_setmaxnreg_repro/Cargo.toml | 12 + .../examples/tma_setmaxnreg_repro/README.md | 18 + .../examples/tma_setmaxnreg_repro/src/main.rs | 93 +++ .../tma_setmaxnreg_repro/verify-code-shape.sh | 71 ++ scripts/smoketest.sh | 21 +- 19 files changed, 1246 insertions(+), 12 deletions(-) create mode 100644 crates/mir-lower/tests/lowering_test/tma.rs create mode 100644 crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/Cargo.lock create mode 100644 crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/Cargo.toml create mode 100644 crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/README.md create mode 100644 crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/src/main.rs create mode 100755 crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/verify-code-shape.sh diff --git a/Cargo.lock b/Cargo.lock index f3a0500727..d97b92e22c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1011,6 +1011,7 @@ version = "0.2.1" dependencies = [ "anyhow", "combine", + "cuda-target-spec", "dialect-mir", "dialect-nvvm", "llvm-export", diff --git a/crates/cuda-oxide-codegen/src/iket.rs b/crates/cuda-oxide-codegen/src/iket.rs index cf48440d1e..a82c1ca477 100644 --- a/crates/cuda-oxide-codegen/src/iket.rs +++ b/crates/cuda-oxide-codegen/src/iket.rs @@ -942,6 +942,7 @@ mod tests { module, true, mir_lower::IntrinsicBackend::LlvmNvptx, + None, ) .unwrap(); assert!(!has_iket_operations(&ctx, module)); diff --git a/crates/cuda-oxide-codegen/src/lower.rs b/crates/cuda-oxide-codegen/src/lower.rs index 83995f95cf..26ca03230e 100644 --- a/crates/cuda-oxide-codegen/src/lower.rs +++ b/crates/cuda-oxide-codegen/src/lower.rs @@ -35,6 +35,7 @@ pub fn lower_to_llvm( module_op_ptr: Ptr, allow_fma_contraction: bool, intrinsic_backend: mir_lower::IntrinsicBackend, + target_arch: Option, ) -> Result<(), PipelineError> { mir_lower::register(ctx); @@ -44,6 +45,7 @@ pub fn lower_to_llvm( mir_lower::LoweringOptions { allow_fma_contraction, intrinsic_backend, + target_arch, }, ) { Ok(()) => Ok(()), diff --git a/crates/cuda-oxide-codegen/src/pipeline.rs b/crates/cuda-oxide-codegen/src/pipeline.rs index 68786bc411..98201360b9 100644 --- a/crates/cuda-oxide-codegen/src/pipeline.rs +++ b/crates/cuda-oxide-codegen/src/pipeline.rs @@ -30,9 +30,11 @@ use crate::ptx::{ }; use crate::target::detect_features_in_llvm_text; use crate::verify::verify_operation; +use cuda_target_spec::CudaArch; use llvm_export::export::{DebugKind, FunctionLocalStaticPlacement, NvvmIrDialect}; use pliron::context::{Context, Ptr}; use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; use pliron::operation::Operation; use pliron::printable::Printable; use std::path::{Path, PathBuf}; @@ -194,25 +196,28 @@ pub fn compile_translated_module( strip_iket(ctx, module)?; } - // IKET's placeholder ABI is keyed by the concrete sm_* family, but the - // definitive target is normally resolved only after LLVM lowering + // IKET's placeholder ABI and the SM120/SM121 local TMA G2S spelling are + // keyed by the concrete sm_* family, but the definitive target is normally + // resolved only after LLVM lowering // (`generate_ptx_discovered` on the PTX path, `resolve_nvvm_target_with_generated` // on the NVVM path), where a device hint that cannot lower a detected // feature is silently raised to the feature floor. Materializing from // the pre-resolution hint could then bake a placeholder shape for a - // family the module is never compiled for. So when IKET operations are - // present, promote the hint to the pipeline's explicit target: both - // resolvers honor an explicit target exactly (they validate it and fail + // family the module is never compiled for. So when target-dependent + // lowering is present, promote the hint to the pipeline's explicit target: + // both resolvers honor an explicit target exactly (they validate it and fail // loudly instead of raising it), so the placeholder shape and the // compiled target can no longer diverge. + let lowering_needs_concrete_target = + has_iket_operations(ctx, module) || has_unicast_tma_g2s(ctx, module); let pinned_backend: BackendOptions; let backend: &BackendOptions = if request.backend.target_arch.is_none() && request.backend.device_arch_hint.is_some() - && has_iket_operations(ctx, module) + && lowering_needs_concrete_target { pinned_backend = BackendOptions { target_arch: request.backend.device_arch_hint.clone(), - target_arch_source: "the detected GPU, pinned by IKET materialization", + target_arch_source: "the detected GPU, pinned before target-dependent MIR lowering", ..request.backend.clone() }; &pinned_backend @@ -318,11 +323,24 @@ pub fn compile_translated_module( .trace .emit("\n=== Lowering dialect-mir → LLVM dialect ==="); } + let lowering_target = backend + .target_arch + .as_deref() + .map(|target| { + target + .parse::() + .map_err(|error| PipelineError::TargetSelection { + target: target.to_owned(), + reason: format!("{error} (target from {})", backend.target_arch_source), + }) + }) + .transpose()?; lower_to_llvm( ctx, module, !backend.no_fma, backend_selection.intrinsic_backend, + lowering_target, )?; let lowered_module_uses_libdevice = module_uses_libdevice(ctx, module); @@ -590,6 +608,31 @@ pub fn compile_translated_module( }) } +fn has_unicast_tma_g2s(ctx: &Context, operation: Ptr) -> bool { + use dialect_nvvm::ops::{ + CpAsyncBulkTensorG2sTile1dOp, CpAsyncBulkTensorG2sTile2dOp, CpAsyncBulkTensorG2sTile3dOp, + CpAsyncBulkTensorG2sTile4dOp, CpAsyncBulkTensorG2sTile5dOp, + }; + + let opid = Operation::get_opid(operation, ctx); + if opid == CpAsyncBulkTensorG2sTile1dOp::get_opid_static() + || opid == CpAsyncBulkTensorG2sTile2dOp::get_opid_static() + || opid == CpAsyncBulkTensorG2sTile3dOp::get_opid_static() + || opid == CpAsyncBulkTensorG2sTile4dOp::get_opid_static() + || opid == CpAsyncBulkTensorG2sTile5dOp::get_opid_static() + { + return true; + } + operation.deref(ctx).regions().any(|region| { + region.deref(ctx).iter(ctx).any(|block| { + block + .deref(ctx) + .iter(ctx) + .any(|child| has_unicast_tma_g2s(ctx, child)) + }) + }) +} + /// Backend decision made from the typed module before MIR lowering starts. #[derive(Clone, Copy, Debug, PartialEq, Eq)] struct PreLoweringBackendSelection { diff --git a/crates/mir-lower/Cargo.toml b/crates/mir-lower/Cargo.toml index d2fa417eb1..5fdd847e50 100644 --- a/crates/mir-lower/Cargo.toml +++ b/crates/mir-lower/Cargo.toml @@ -9,6 +9,7 @@ description = "dialect-mir to LLVM dialect lowering pass" readme = "README.md" [dependencies] +cuda-target-spec = { workspace = true } rustc-hash = { workspace = true } pliron = { workspace = true } llvm-export = { workspace = true } diff --git a/crates/mir-lower/src/context.rs b/crates/mir-lower/src/context.rs index 290e3d5073..bf1bdf1536 100644 --- a/crates/mir-lower/src/context.rs +++ b/crates/mir-lower/src/context.rs @@ -45,7 +45,7 @@ pub(crate) fn lowering_options(ctx: &Context) -> LoweringOptions { ctx.aux_data_map .get(&*options_storage::LOWERING_OPTIONS_KEY) .and_then(|index| ctx.aux_data[*index].downcast_ref::()) - .copied() + .cloned() .unwrap_or_default() } diff --git a/crates/mir-lower/src/convert/intrinsics/tma.rs b/crates/mir-lower/src/convert/intrinsics/tma.rs index 32984af18c..a39545eb38 100644 --- a/crates/mir-lower/src/convert/intrinsics/tma.rs +++ b/crates/mir-lower/src/convert/intrinsics/tma.rs @@ -8,6 +8,9 @@ use crate::convert::intrinsics::common::*; use crate::helpers; use crate::{IntrinsicBackend, context}; +use cuda_target_spec::CudaArch; +use dialect_mir::types::address_space; +use llvm_export::op_interfaces::CastOpInterface; use llvm_export::ops as llvm; use llvm_export::types as llvm_types; use pliron::builtin::op_interfaces::CallOpCallable; @@ -19,6 +22,7 @@ use pliron::irbuild::rewriter::Rewriter; use pliron::op::Op; use pliron::operation::Operation; use pliron::result::Result; +use pliron::r#type::Typed; /// Convert TMA G2S (global to shared) operations using LLVM intrinsics. pub(crate) fn convert_g2s( @@ -56,6 +60,68 @@ fn g2s_inline_asm(dims: usize, multicast: bool, cta_group: i32) -> (String, Stri (template, constraints.join(",")) } +fn g2s_cta_inline_asm(dims: usize) -> (String, String) { + let coordinates = (0..dims) + .map(|index| format!("${}", 3 + index)) + .collect::>() + .join(", "); + let template = format!( + "cp.async.bulk.tensor.{dims}d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$0], [$2, {{{coordinates}}}], [$1];" + ); + let mut constraints = vec!["r", "r", "l"]; + constraints.extend(std::iter::repeat_n("r", dims)); + constraints.push("~{memory}"); + (template, constraints.join(",")) +} + +fn target_requires_cta_local_g2s(target: Option<&CudaArch>) -> bool { + target.is_some_and(|target| matches!(target.capability(), 120 | 121)) +} + +fn pointer_address_space(ctx: &Context, value: pliron::value::Value) -> Result { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .map(llvm_types::PointerType::address_space) + .ok_or_else(|| pliron::input_error_noloc!("TMA G2S destination must be a pointer")) +} + +fn lower_cta_local_g2s( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + operands: &[pliron::value::Value], + dims: usize, +) { + let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); + let dst = cast_to_shared_addrspace(ctx, rewriter, operands[0]); + let dst_address = llvm::PtrToIntOp::new(ctx, dst, i32_ty.into()); + rewriter.insert_operation(ctx, dst_address.get_operation()); + let barrier = cast_to_shared_addrspace(ctx, rewriter, operands[1]); + let barrier_address = llvm::PtrToIntOp::new(ctx, barrier, i32_ty.into()); + rewriter.insert_operation(ctx, barrier_address.get_operation()); + + let mut inputs = vec![ + dst_address.get_operation().deref(ctx).get_result(0), + barrier_address.get_operation().deref(ctx).get_result(0), + operands[2], + ]; + inputs.extend(operands[3..3 + dims].iter().copied()); + let (template, constraints) = g2s_cta_inline_asm(dims); + let void_ty = llvm_types::VoidType::get(ctx); + inline_asm_convergent( + ctx, + rewriter, + op, + void_ty.into(), + inputs, + &template, + &constraints, + ); + rewriter.erase_operation(ctx, op); +} + fn convert_g2s_impl( ctx: &mut Context, rewriter: &mut DialectConversionRewriter, @@ -85,10 +151,22 @@ fn convert_g2s_impl( ); } + let options = context::lowering_options(ctx); + if !multicast && target_requires_cta_local_g2s(options.target_arch.as_ref()) { + if pointer_address_space(ctx, operands[0])? == address_space::CLUSTER_SHARED { + return pliron::input_err_noloc!( + "TMA G2S on {} cannot target cluster-shared address space 7", + options.target_arch.as_ref().unwrap() + ); + } + lower_cta_local_g2s(ctx, rewriter, op, &operands, dims); + return Ok(()); + } + let dst_casted = cast_to_cluster_shared_addrspace(ctx, rewriter, operands[0]); let barrier_casted = cast_to_shared_addrspace(ctx, rewriter, operands[1]); - if context::lowering_options(ctx).intrinsic_backend == IntrinsicBackend::LibNvvm { + if options.intrinsic_backend == IntrinsicBackend::LibNvvm { let mut inputs = vec![dst_casted, barrier_casted, operands[2]]; inputs.extend(operands[3..3 + dims].iter().copied()); if multicast { @@ -722,7 +800,7 @@ pub(crate) fn convert_control( #[cfg(test)] mod tests { - use super::{g2s_inline_asm, reduce_inline_asm, s2g_inline_asm}; + use super::{g2s_cta_inline_asm, g2s_inline_asm, reduce_inline_asm, s2g_inline_asm}; #[test] fn inline_tma_templates_keep_exact_ptx_shapes() { @@ -740,6 +818,13 @@ mod tests { "l,l,l,r,r,h,~{memory}".into(), ) ); + assert_eq!( + g2s_cta_inline_asm(5), + ( + "cp.async.bulk.tensor.5d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$0], [$2, {$3, $4, $5, $6, $7}], [$1];".into(), + "r,r,l,r,r,r,r,r,~{memory}".into(), + ) + ); assert_eq!( s2g_inline_asm(5), ( diff --git a/crates/mir-lower/src/convert/ops/arithmetic.rs b/crates/mir-lower/src/convert/ops/arithmetic.rs index b57f7e1d6b..5fd415f0c1 100644 --- a/crates/mir-lower/src/convert/ops/arithmetic.rs +++ b/crates/mir-lower/src/convert/ops/arithmetic.rs @@ -995,6 +995,7 @@ mod tests { crate::LoweringOptions { allow_fma_contraction: false, intrinsic_backend: crate::IntrinsicBackend::LlvmNvptx, + ..Default::default() }, ) .expect("lowering failed"); diff --git a/crates/mir-lower/src/convert/ops/call.rs b/crates/mir-lower/src/convert/ops/call.rs index 93bd1bdc9e..43c61dff55 100644 --- a/crates/mir-lower/src/convert/ops/call.rs +++ b/crates/mir-lower/src/convert/ops/call.rs @@ -2014,6 +2014,7 @@ mod tests { crate::LoweringOptions { allow_fma_contraction: false, intrinsic_backend: crate::IntrinsicBackend::LlvmNvptx, + ..Default::default() }, ); @@ -2502,6 +2503,7 @@ mod tests { crate::LoweringOptions { allow_fma_contraction: true, intrinsic_backend: crate::IntrinsicBackend::LlvmNvptx, + ..Default::default() }, ); assert_eq!( @@ -2550,6 +2552,7 @@ mod tests { crate::LoweringOptions { allow_fma_contraction: true, intrinsic_backend: crate::IntrinsicBackend::LibNvvm, + ..Default::default() }, ); assert_eq!( diff --git a/crates/mir-lower/src/lib.rs b/crates/mir-lower/src/lib.rs index a93407420a..d511792771 100644 --- a/crates/mir-lower/src/lib.rs +++ b/crates/mir-lower/src/lib.rs @@ -175,7 +175,7 @@ pub enum IntrinsicBackend { } /// Options controlling the `dialect-mir` to LLVM dialect lowering pass. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq)] pub struct LoweringOptions { /// Whether ordinary floating-point multiply/add or multiply/subtract /// expressions may contract into fused operations. @@ -184,6 +184,9 @@ pub struct LoweringOptions { pub allow_fma_contraction: bool, /// Intrinsic ABI expected by the selected LLVM-to-device backend. pub intrinsic_backend: IntrinsicBackend, + /// Concrete CUDA target when lowering depends on an architecture-specific + /// PTX form. `None` keeps target-independent historical lowering. + pub target_arch: Option, } impl Default for LoweringOptions { @@ -191,6 +194,7 @@ impl Default for LoweringOptions { Self { allow_fma_contraction: true, intrinsic_backend: IntrinsicBackend::LlvmNvptx, + target_arch: None, } } } diff --git a/crates/mir-lower/tests/lowering_test/main.rs b/crates/mir-lower/tests/lowering_test/main.rs index cdc0a69728..ee27dca03e 100644 --- a/crates/mir-lower/tests/lowering_test/main.rs +++ b/crates/mir-lower/tests/lowering_test/main.rs @@ -17,5 +17,6 @@ mod math_conversions; mod matrix_memory; mod mma; mod sregs_and_warp; +mod tma; mod wgmma_lowering; mod wgmma_rejections; diff --git a/crates/mir-lower/tests/lowering_test/tma.rs b/crates/mir-lower/tests/lowering_test/tma.rs new file mode 100644 index 0000000000..b319f47391 --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/tma.rs @@ -0,0 +1,145 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::types::{MirPtrType, address_space}; +use dialect_nvvm::ops as nvvm; +use pliron::builtin::ops::ModuleOp; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::printable::Printable; + +use crate::common::{append_return, build_test_kernel, make_test_ctx}; + +fn g2s_module(destination_address_space: u32) -> (Context, Ptr) { + let mut ctx = make_test_ctx(); + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let i16_ty = IntegerType::get(&ctx, 16, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let destination = MirPtrType::get(&mut ctx, u8_ty.into(), true, destination_address_space); + let barrier = MirPtrType::get_shared(&mut ctx, u8_ty.into(), true); + let descriptor = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); + let (module, entry) = build_test_kernel( + &mut ctx, + vec![ + destination.into(), + barrier.into(), + descriptor.into(), + i32_ty.into(), + i32_ty.into(), + i16_ty.into(), + i64_ty.into(), + ], + ); + let operands = (0..7) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect(); + Operation::new( + &mut ctx, + nvvm::CpAsyncBulkTensorG2sTile2dOp::get_concrete_op_info(), + vec![], + operands, + vec![], + 0, + ) + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + (ctx, module) +} + +fn lower_g2s( + destination_address_space: u32, + target: &str, + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, Ptr), String> { + let (mut ctx, module) = g2s_module(destination_address_space); + let result = mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + target_arch: Some(target.parse().unwrap()), + ..Default::default() + }, + ); + match result { + Ok(()) => Ok((ctx, module)), + Err(error) => Err(error.disp(&ctx).to_string()), + } +} + +fn exported_module(ctx: &Context, module: Ptr) -> String { + let module = Operation::get_op::(module, ctx).unwrap(); + llvm_export::export::export_module_to_string(ctx, &module).unwrap() +} + +#[test] +fn sm120_local_g2s_uses_cta_inline_ptx_on_both_backends() { + for target in [ + "sm_120", "sm_120a", "sm_120f", "sm_121", "sm_121a", "sm_121f", + ] { + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + let (ctx, module) = lower_g2s(address_space::SHARED, target, backend).unwrap(); + let ir = exported_module(&ctx, module); + assert!( + ir.contains( + "cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx::bytes" + ), + "target {target}: {ir}" + ); + assert!( + ir.contains("\"r,r,l,r,r,~{memory}\""), + "target {target}: {ir}" + ); + assert!(!ir.contains("shared::cluster"), "target {target}: {ir}"); + assert!( + !ir.contains("llvm.nvvm.cp.async.bulk.tensor.g2s"), + "target {target}: {ir}" + ); + } + } +} + +#[test] +fn sm90_local_g2s_keeps_cluster_intrinsic_lowering() { + let (ctx, module) = lower_g2s( + address_space::SHARED, + "sm_90a", + mir_lower::IntrinsicBackend::LlvmNvptx, + ) + .unwrap(); + let ir = exported_module(&ctx, module); + assert!( + ir.contains("llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d"), + "{ir}" + ); + assert!( + ir.contains("addrspacecast ptr addrspace(3)") && ir.contains("to ptr addrspace(7)"), + "{ir}" + ); + assert!(!ir.contains("shared::cta"), "{ir}"); +} + +#[test] +fn sm120_rejects_cluster_shared_destination_for_unicast_g2s() { + let result = lower_g2s( + address_space::CLUSTER_SHARED, + "sm_120a", + mir_lower::IntrinsicBackend::LlvmNvptx, + ); + let error = match result { + Ok(_) => panic!("SM120 unicast G2S accepted an AS7 destination"), + Err(error) => error, + }; + assert!( + error.contains("cannot target cluster-shared address space 7"), + "{error}" + ); +} diff --git a/crates/rustc-codegen-cuda/Cargo.lock b/crates/rustc-codegen-cuda/Cargo.lock index d7542a35bf..0a76350ca1 100644 --- a/crates/rustc-codegen-cuda/Cargo.lock +++ b/crates/rustc-codegen-cuda/Cargo.lock @@ -517,6 +517,7 @@ name = "mir-lower" version = "0.2.1" dependencies = [ "anyhow", + "cuda-target-spec", "dialect-mir", "dialect-nvvm", "llvm-export", diff --git a/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/Cargo.lock b/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/Cargo.lock new file mode 100644 index 0000000000..9f386b2bf7 --- /dev/null +++ b/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/Cargo.lock @@ -0,0 +1,733 @@ +# 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.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.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8498c871161e1742aaa9d52551b2d6ebdd4c3d45a3be423e3728f33b955be550" +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.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.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84bdb592a2844148fb13bdfbbab158879f56916bde1d1a1b6580ca27f23850cd" +dependencies = [ + "anyhow", + "cuda-bindings", + "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" +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 = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +dependencies = [ + "libloading", + "thiserror", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "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.5", +] + +[[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 = "tma_setmaxnreg_repro" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.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/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/Cargo.toml b/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/Cargo.toml new file mode 100644 index 0000000000..c75c7e6fe8 --- /dev/null +++ b/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "tma_setmaxnreg_repro" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" + +[workspace] + +[dependencies] +cuda-core = "0.3.1" +cuda-device = { path = "../../../cuda-device" } +cuda-host = { path = "../../../cuda-host" } diff --git a/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/README.md b/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/README.md new file mode 100644 index 0000000000..25a5e547b3 --- /dev/null +++ b/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/README.md @@ -0,0 +1,18 @@ +# TMA G2S with `setmaxnreg` + +This compile-only regression contains a 384-thread kernel that combines the +existing `cp_async_bulk_tensor_2d_g2s` API with a 40/232-register warpgroup +split. On SM120, the local shared destination must lower to `shared::cta` so +ptxas can retain both `USETMAXREG` instructions without an address-compatibility +call. A second kernel covers the same lowering for all five TMA dimensions. + +The lowering rule covers capabilities 120 and 121 independently of the target +suffix, including `sm_120f` and `sm_121a`. The smoketest builds both compiler +routes at `sm_120a`. Its code-shape check +requires CTA-local PTX on the LLVM path, inline CTA-local assembly in libNVVM +IR, and SASS containing `UTMALDG.2D` plus the 40/232 register reallocation with +no `CALL.ABS`. + +```sh +bash scripts/smoketest.sh --only '^tma_setmaxnreg_repro$' +``` diff --git a/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/src/main.rs b/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/src/main.rs new file mode 100644 index 0000000000..05f6088fe0 --- /dev/null +++ b/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/src/main.rs @@ -0,0 +1,93 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-only regression for SM120 TMA G2S with warpgroup register reallocation. + +use cuda_device::{ + SharedArray, + barrier::{Barrier, mbarrier_init}, + kernel, launch_bounds, thread, + tma::{ + TmaDescriptor, cp_async_bulk_tensor_1d_g2s, cp_async_bulk_tensor_2d_g2s, + cp_async_bulk_tensor_3d_g2s, cp_async_bulk_tensor_4d_g2s, cp_async_bulk_tensor_5d_g2s, + }, +}; +use cuda_host::cuda_module; + +const THREADS: u32 = 384; +const MATH_THREADS: u32 = 256; + +#[cuda_module] +mod kernels { + use super::*; + + #[kernel] + #[launch_bounds(THREADS, 1)] + pub unsafe fn setmaxnreg_control(output: *mut u32) { + let thread = thread::threadIdx_x(); + if thread >= MATH_THREADS { + unsafe { thread::setmaxnreg_dec::<40>() }; + return; + } + unsafe { thread::setmaxnreg_inc::<232>() }; + if thread == 0 { + unsafe { output.write(1) }; + } + } + + #[kernel] + #[launch_bounds(THREADS, 1)] + pub unsafe fn setmaxnreg_with_tma(tensor_map: *const TmaDescriptor, output: *mut u32) { + static mut TILE: SharedArray = SharedArray::UNINIT; + static mut READY: Barrier = Barrier::UNINIT; + + let thread = thread::threadIdx_x(); + if thread == 0 { + unsafe { mbarrier_init(&raw mut READY, 1) }; + } + thread::sync_threads(); + if thread >= MATH_THREADS { + unsafe { thread::setmaxnreg_dec::<40>() }; + if thread == MATH_THREADS { + unsafe { + cp_async_bulk_tensor_2d_g2s( + (&raw mut TILE).cast(), + tensor_map, + 0, + 0, + &raw mut READY, + ); + } + } + return; + } + unsafe { thread::setmaxnreg_inc::<232>() }; + if thread == 0 { + unsafe { output.write(1) }; + } + } + + #[kernel] + #[launch_bounds(32, 1)] + pub unsafe fn g2s_all_dimensions(tensor_map: *const TmaDescriptor) { + static mut TILE: SharedArray = SharedArray::UNINIT; + static mut READY: Barrier = Barrier::UNINIT; + + if thread::threadIdx_x() == 0 { + unsafe { + mbarrier_init(&raw mut READY, 1); + let tile = (&raw mut TILE).cast(); + let barrier = &raw mut READY; + cp_async_bulk_tensor_1d_g2s(tile, tensor_map, 0, barrier); + cp_async_bulk_tensor_2d_g2s(tile, tensor_map, 0, 0, barrier); + cp_async_bulk_tensor_3d_g2s(tile, tensor_map, 0, 0, 0, barrier); + cp_async_bulk_tensor_4d_g2s(tile, tensor_map, 0, 0, 0, 0, barrier); + cp_async_bulk_tensor_5d_g2s(tile, tensor_map, 0, 0, 0, 0, 0, barrier); + } + } + } +} + +fn main() {} diff --git a/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/verify-code-shape.sh b/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/verify-code-shape.sh new file mode 100755 index 0000000000..b8ee0c82c1 --- /dev/null +++ b/crates/rustc-codegen-cuda/examples/tma_setmaxnreg_repro/verify-code-shape.sh @@ -0,0 +1,71 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +root="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +mode="${1:-ptx}" +target="${2:-sm_120a}" +cta='cp\.async\.bulk\.tensor\.[1-5]d\.shared::cta\.global\.tile\.mbarrier::complete_tx::bytes' +cluster='cp\.async\.bulk\.tensor\.[1-5]d\.shared::cluster\.global\.tile\.mbarrier::complete_tx::bytes' + +case "${mode}" in + ptx) + artifact="${root}/tma_setmaxnreg_repro.ptx" + test -s "${artifact}" + grep -qx "\.target ${target}" "${artifact}" + [[ "$(grep -Ec "${cta}" "${artifact}")" -eq 6 ]] + [[ "$(grep -Eo "${cta}" "${artifact}" | sort -u | wc -l)" -eq 5 ]] + ! grep -Eq "${cluster}" "${artifact}" + + ptxas_bin="" + for candidate in "${CUDA_OXIDE_PTXAS:-}" \ + "$(command -v ptxas 2>/dev/null)" \ + "${CUDA_HOME:+${CUDA_HOME}/bin/ptxas}" \ + /usr/local/cuda/bin/ptxas \ + /usr/local/cuda-*/bin/ptxas; do + if [[ -n "${candidate}" && -x "${candidate}" ]]; then + ptxas_bin="${candidate}" + break + fi + done + if [[ -z "${ptxas_bin}" ]]; then + echo "error: ptxas is required for the setmaxnreg SASS regression" >&2 + exit 1 + fi + nvdisasm_bin="${CUDA_OXIDE_NVDISASM:-${ptxas_bin%/ptxas}/nvdisasm}" + if [[ ! -x "${nvdisasm_bin}" ]]; then + nvdisasm_bin="$(command -v nvdisasm 2>/dev/null)" + fi + if [[ -z "${nvdisasm_bin}" || ! -x "${nvdisasm_bin}" ]]; then + echo "error: nvdisasm is required for the setmaxnreg SASS regression" >&2 + exit 1 + fi + + scratch="$(mktemp -d /tmp/cuda-oxide-tma-setmaxnreg.XXXXXX)" + trap 'rm -rf "${scratch}"' EXIT + "${ptxas_bin}" --gpu-name="${target}" --verbose "${artifact}" \ + --output-file="${scratch}/repro.cubin" 2>"${scratch}/ptxas.log" + ! grep -q 'C7506' "${scratch}/ptxas.log" + "${nvdisasm_bin}" --print-code "${scratch}/repro.cubin" >"${scratch}/sass.txt" + [[ "$(grep -c 'USETMAXREG.DEALLOC' "${scratch}/sass.txt")" -eq 2 ]] + [[ "$(grep -c 'USETMAXREG.TRY_ALLOC' "${scratch}/sass.txt")" -eq 2 ]] + grep -q 'UTMALDG.2D' "${scratch}/sass.txt" + ! grep -q 'CALL.ABS' "${scratch}/sass.txt" + ;; + nvvm) + artifact="${root}/tma_setmaxnreg_repro.ll" + test -s "${artifact}" + [[ "$(grep -Ec "${cta}" "${artifact}")" -eq 6 ]] + [[ "$(grep -Eo "${cta}" "${artifact}" | sort -u | wc -l)" -eq 5 ]] + ! grep -Eq "${cluster}" "${artifact}" + ! grep -q 'llvm\.nvvm\.cp\.async\.bulk\.tensor\.g2s\.cta' "${artifact}" + ;; + *) + echo "usage: $0 [ptx|nvvm] [sm_target]" >&2 + exit 2 + ;; +esac + +echo "tma_setmaxnreg_repro ${mode} code shape: PASS" diff --git a/scripts/smoketest.sh b/scripts/smoketest.sh index 36811ddc27..54759ef5ce 100755 --- a/scripts/smoketest.sh +++ b/scripts/smoketest.sh @@ -73,7 +73,7 @@ LTOIR_EXAMPLES=(addressof_sharedarray cpp_consumes_rust_device device_ffi_test l LTOIR_MODERN_EXAMPLES=(small_type_ffi_test) AUTO_NVVM_EXAMPLES=(libdevice_math) IKET_EXAMPLES=(iket_trace) -BLACKWELL_COMPILE_EXAMPLES=(generated_intrinsics_blackwell) +BLACKWELL_COMPILE_EXAMPLES=(generated_intrinsics_blackwell tma_setmaxnreg_repro) SM100_COMPILE_EXAMPLES=(redux_f32) NVVM_VERIFY_EXAMPLES=(cp_async_small device_global enum_constant_provenance ex2_approx_f16 generated_intrinsics generated_intrinsics_blackwell generated_ldmatrix legacy_atomic_fadd legacy_atomic_rmw_cas libdevice_math legacy_nvvm_pointer_shapes packed_atomic_add primitive_stress scoped_atomic_load_store shuffle_64 tcgen05 tuple_constant_provenance wgmma_mma_bf16) ERROR_EXAMPLES=(error error_set_discriminant_uninhabited error_enum_bool_payload_addr error_enum_pointer_overlap error_enum_shared_pointer_layout error_heap_alloc error_host_arch_intrinsic error_host_target_feature error_kernel_shared_param error_missing_device_attr error_generated_intrinsic_abi error_generated_intrinsic_unknown_id error_generated_intrinsic_fn_pointer error_generated_intrinsic_callable) @@ -925,6 +925,25 @@ run_cargo() { CARGO_EC=${llvm_ec} return fi + if [[ "${ex}" == "tma_setmaxnreg_repro" ]]; then + local shape_check="crates/rustc-codegen-cuda/examples/${ex}/verify-code-shape.sh" + if ! bash "${shape_check}" ptx >>"${log}" 2>&1; then + CARGO_EC=1 + return + fi + if [[ ${VERBOSE} -eq 1 ]]; then + cargo oxide "${nvvm_args[@]}" 2>&1 | tee -a "${log}" + CARGO_EC=${PIPESTATUS[0]} + else + cargo oxide "${nvvm_args[@]}" >>"${log}" 2>&1 + CARGO_EC=$? + fi + if [[ ${CARGO_EC} -eq 0 ]] \ + && ! bash "${shape_check}" nvvm >>"${log}" 2>&1; then + CARGO_EC=1 + fi + return + fi local llvm_ptx="crates/rustc-codegen-cuda/examples/${ex}/${ex}.ptx" local instruction_re='mma\.sp::ordered_metadata\.sync\.aligned\.m16n8k64\.row\.col\.kind::f8f6f4\.f32\.[[:alnum:]]+\.[[:alnum:]]+\.f32' local sparse_f16_instruction_re='mma\.sp::ordered_metadata\.sync\.aligned\.m16n8k64\.row\.col\.kind::f8f6f4\.f16\.(e2m1|e2m3|e3m2|e4m3|e5m2)\.(e2m1|e2m3|e3m2|e4m3|e5m2)\.f16'