fix(tma): select CTA-local G2S on SM120/SM121

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));

crates/cuda-oxide-codegen/src/lower.rs

@@ -35,6 +35,7 @@ pub fn lower_to_llvm(
     module_op_ptr: Ptr<Operation>,
     allow_fma_contraction: bool,
     intrinsic_backend: mir_lower::IntrinsicBackend,
+    target_arch: Option<cuda_target_spec::CudaArch>,
 ) -> 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(()),

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::<CudaArch>()
+                .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<Operation>) -> 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 {

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 }

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::<LoweringOptions>())
-        .copied()
+        .cloned()
         .unwrap_or_default()
 }
 

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::<Vec<_>>()
+        .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<u32> {
+    value
+        .get_type(ctx)
+        .deref(ctx)
+        .downcast_ref::<llvm_types::PointerType>()
+        .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<Operation>,
+    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),
             (

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");

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!(

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,13 +184,17 @@ 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<cuda_target_spec::CudaArch>,
 }
 
 impl Default for LoweringOptions {
     fn default() -> Self {
         Self {
             allow_fma_contraction: true,
             intrinsic_backend: IntrinsicBackend::LlvmNvptx,
+            target_arch: None,
         }
     }
 }

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;